;; ==========================================================================
;; types.sx — Gradual type system for SX
;;
;; Registration-time type checking: zero runtime cost.
;; Annotations are optional — unannotated code defaults to `any`.
;;
;; Depends on: eval.sx (type-of, component accessors, env ops)
;;             primitives.sx, boundary.sx (return type declarations)
;;
;; Platform interface (from eval.sx, already provided):
;;   (type-of x)            → type string
;;   (symbol-name s)        → string
;;   (keyword-name k)       → string
;;   (component-body c)     → AST
;;   (component-params c)   → list of param name strings
;;   (component-has-children c) → boolean
;;   (env-get env k)        → value or nil
;;
;; New platform functions for types.sx:
;;   (component-param-types c)  → dict {param-name → type} or nil
;;   (component-set-param-types! c d) → store param types on component
;; ==========================================================================


;; --------------------------------------------------------------------------
;; 1. Type representation
;; --------------------------------------------------------------------------
;; Types are plain SX values:
;;   - Strings for base types: "number", "string", "boolean", "nil",
;;     "symbol", "keyword", "element", "any", "never"
;;   - Nullable shorthand: "string?", "number?", "dict?", "boolean?"
;;     → equivalent to (or string nil) etc.
;;   - Lists for compound types:
;;     (or t1 t2 ...)         — union
;;     (list-of t)            — homogeneous list
;;     (dict-of tk tv)        — typed dict
;;     (-> t1 t2 ... treturn) — function type (last is return)

;; Base type names
(define base-types
  (list "number" "string" "boolean" "nil" "symbol" "keyword"
        "element" "any" "never" "list" "dict"
        "lambda" "component" "island" "macro" "signal"))


;; --------------------------------------------------------------------------
;; 2. Type predicates
;; --------------------------------------------------------------------------

(define type-any?
  (fn (t) (= t "any")))

(define type-never?
  (fn (t) (= t "never")))

(define type-nullable?
  (fn (t)
    ;; A type is nullable if it's "any", "nil", a "?" shorthand, or
    ;; a union containing "nil".
    (if (= t "any") true
      (if (= t "nil") true
        (if (and (= (type-of t) "string") (ends-with? t "?")) true
          (if (and (= (type-of t) "list")
                   (not (empty? t))
                   (= (first t) "or"))
            (contains? (rest t) "nil")
            false))))))

(define nullable-base
  (fn (t)
    ;; Strip "?" from nullable shorthand: "string?" → "string"
    (if (and (= (type-of t) "string")
             (ends-with? t "?")
             (not (= t "?")))
      (slice t 0 (- (string-length t) 1))
      t)))


;; --------------------------------------------------------------------------
;; 3. Subtype checking
;; --------------------------------------------------------------------------
;; subtype?(a, b) — is type `a` assignable to type `b`?

(define subtype?
  (fn (a b)
    ;; any accepts everything
    (if (type-any? b) true
      ;; never is subtype of everything
      (if (type-never? a) true
        ;; any is not a subtype of a specific type
        (if (type-any? a) false
          ;; identical types
          (if (= a b) true
            ;; nil is subtype of nullable types
            (if (= a "nil")
              (type-nullable? b)
              ;; nullable shorthand: "string?" = (or string nil)
              (if (and (= (type-of b) "string") (ends-with? b "?"))
                (let ((base (nullable-base b)))
                  (or (= a base) (= a "nil")))
                ;; a is a union: (or t1 t2 ...) <: b if ALL members <: b
                ;; Must check before b-union — (or A B) <: (or A B C) needs
                ;; each member of a checked against the full union b.
                (if (and (= (type-of a) "list")
                         (not (empty? a))
                         (= (first a) "or"))
                  (every? (fn (member) (subtype? member b)) (rest a))
                  ;; union: a <: (or t1 t2 ...) if a <: any member
                  (if (and (= (type-of b) "list")
                           (not (empty? b))
                           (= (first b) "or"))
                    (some (fn (member) (subtype? a member)) (rest b))
                    ;; list-of covariance
                    (if (and (= (type-of a) "list") (= (type-of b) "list")
                             (= (len a) 2) (= (len b) 2)
                             (= (first a) "list-of") (= (first b) "list-of"))
                      (subtype? (nth a 1) (nth b 1))
                      ;; "list" <: (list-of any)
                      (if (and (= a "list")
                               (= (type-of b) "list")
                               (= (len b) 2)
                               (= (first b) "list-of"))
                        (type-any? (nth b 1))
                        ;; (list-of t) <: "list"
                        (if (and (= (type-of a) "list")
                                 (= (len a) 2)
                                 (= (first a) "list-of")
                                 (= b "list"))
                          true
                          ;; "element" is subtype of "string?" (rendered HTML)
                          false)))))))))))))


;; --------------------------------------------------------------------------
;; 4. Type union
;; --------------------------------------------------------------------------

(define type-union
  (fn (a b)
    ;; Compute the smallest type that encompasses both a and b.
    (if (= a b) a
      (if (type-any? a) "any"
        (if (type-any? b) "any"
          (if (type-never? a) b
            (if (type-never? b) a
              (if (subtype? a b) b
                (if (subtype? b a) a
                  ;; neither is subtype — create a union
                  (if (= a "nil")
                    ;; nil + string → string?
                    (if (and (= (type-of b) "string")
                             (not (ends-with? b "?")))
                      (str b "?")
                      (list "or" a b))
                    (if (= b "nil")
                      (if (and (= (type-of a) "string")
                               (not (ends-with? a "?")))
                        (str a "?")
                        (list "or" a b))
                      (list "or" a b))))))))))))


;; --------------------------------------------------------------------------
;; 5. Type narrowing
;; --------------------------------------------------------------------------

(define narrow-type
  (fn (t (predicate-name :as string))
    ;; Narrow type based on a predicate test in a truthy branch.
    ;; (if (nil? x) ..then.. ..else..) → in else, x excludes nil.
    ;; Returns (narrowed-then narrowed-else).
    (if (= predicate-name "nil?")
      (list "nil" (narrow-exclude-nil t))
      (if (= predicate-name "string?")
        (list "string" (narrow-exclude t "string"))
        (if (= predicate-name "number?")
          (list "number" (narrow-exclude t "number"))
          (if (= predicate-name "list?")
            (list "list" (narrow-exclude t "list"))
            (if (= predicate-name "dict?")
              (list "dict" (narrow-exclude t "dict"))
              (if (= predicate-name "boolean?")
                (list "boolean" (narrow-exclude t "boolean"))
                ;; Unknown predicate — no narrowing
                (list t t)))))))))


(define narrow-exclude-nil
  (fn (t)
    ;; Remove nil from a type.
    (if (= t "nil") "never"
      (if (= t "any") "any"   ;; can't narrow any
        (if (and (= (type-of t) "string") (ends-with? t "?"))
          (nullable-base t)
          (if (and (= (type-of t) "list")
                   (not (empty? t))
                   (= (first t) "or"))
            (let ((members (filter (fn (m) (not (= m "nil"))) (rest t))))
              (if (= (len members) 1) (first members)
                (if (empty? members) "never"
                  (cons "or" members))))
            t))))))


(define narrow-exclude
  (fn (t excluded)
    ;; Remove a specific type from a union.
    (if (= t excluded) "never"
      (if (= t "any") "any"
        (if (and (= (type-of t) "list")
                 (not (empty? t))
                 (= (first t) "or"))
          (let ((members (filter (fn (m) (not (= m excluded))) (rest t))))
            (if (= (len members) 1) (first members)
              (if (empty? members) "never"
                (cons "or" members))))
          t)))))


;; --------------------------------------------------------------------------
;; 6. Type inference
;; --------------------------------------------------------------------------
;; infer-type walks an AST node and returns its inferred type.
;; type-env is a dict mapping variable names → types.

(define infer-type
  (fn (node (type-env :as dict) (prim-types :as dict) type-registry)
    (let ((kind (type-of node)))
      (if (= kind "number")  "number"
      (if (= kind "string")  "string"
      (if (= kind "boolean") "boolean"
      (if (nil? node)         "nil"
      (if (= kind "keyword") "keyword"
      (if (= kind "symbol")
        (let ((name (symbol-name node)))
          ;; Look up in type env
          (if (has-key? type-env name)
            (get type-env name)
            ;; Builtins
            (if (= name "true") "boolean"
              (if (= name "false") "boolean"
                (if (= name "nil") "nil"
                  ;; Check primitive return types
                  (if (has-key? prim-types name)
                    (get prim-types name)
                    "any"))))))
      (if (= kind "dict") "dict"
      (if (= kind "list")
        (infer-list-type node type-env prim-types type-registry)
        "any")))))))))))


(define infer-list-type
  (fn (node (type-env :as dict) (prim-types :as dict) type-registry)
    ;; Infer type of a list expression (function call, special form, etc.)
    (if (empty? node) "list"
      (let ((head (first node))
            (args (rest node)))
        (if (not (= (type-of head) "symbol"))
          "any"   ;; complex head — can't infer
          (let ((name (symbol-name head)))
            ;; Special forms
            (if (= name "if")
              (infer-if-type args type-env prim-types type-registry)
            (if (= name "when")
              (if (>= (len args) 2)
                (type-union (infer-type (last args) type-env prim-types type-registry) "nil")
                "nil")
            (if (or (= name "cond") (= name "case"))
              "any"   ;; complex — could be refined later
            (if (= name "let")
              (infer-let-type args type-env prim-types type-registry)
            (if (or (= name "do") (= name "begin"))
              (if (empty? args) "nil"
                (infer-type (last args) type-env prim-types type-registry))
            (if (or (= name "lambda") (= name "fn"))
              "lambda"
            (if (= name "and")
              (if (empty? args) "boolean"
                (infer-type (last args) type-env prim-types type-registry))
            (if (= name "or")
              (if (empty? args) "boolean"
                ;; or returns first truthy — union of all args
                (reduce type-union "never"
                  (map (fn (a) (infer-type a type-env prim-types type-registry)) args)))
            (if (= name "map")
              ;; map returns a list
              (if (>= (len args) 2)
                (let ((fn-type (infer-type (first args) type-env prim-types type-registry)))
                  ;; If the fn's return type is known, produce (list-of return-type)
                  (if (and (= (type-of fn-type) "list")
                           (= (first fn-type) "->"))
                    (list "list-of" (last fn-type))
                    "list"))
                "list")
            (if (= name "filter")
              ;; filter preserves element type
              (if (>= (len args) 2)
                (infer-type (nth args 1) type-env prim-types type-registry)
                "list")
            (if (= name "reduce")
              ;; reduce returns the accumulator type — too complex to infer
              "any"
            (if (= name "list")
              "list"
            (if (= name "dict")
              "dict"
            (if (= name "quote")
              "any"
            (if (= name "str")
              "string"
            (if (= name "not")
              "boolean"
            (if (= name "get")
              ;; get — resolve record field type from type registry
              (if (and (>= (len args) 2) (not (nil? type-registry)))
                (let ((dict-type (infer-type (first args) type-env prim-types type-registry))
                      (key-arg (nth args 1))
                      (key-name (cond
                                  (= (type-of key-arg) "keyword") (keyword-name key-arg)
                                  (= (type-of key-arg) "string") key-arg
                                  :else nil)))
                  (if (and key-name
                           (= (type-of dict-type) "string")
                           (has-key? type-registry dict-type))
                    (let ((resolved (resolve-type dict-type type-registry)))
                      (if (and (= (type-of resolved) "dict")
                               (has-key? resolved key-name))
                        (get resolved key-name)
                        "any"))
                    "any"))
                "any")
            (if (starts-with? name "~")
              "element"   ;; component call
              ;; Regular function call: look up return type
              (if (has-key? prim-types name)
                (get prim-types name)
                "any")))))))))))))))))))))))))


(define infer-if-type
  (fn ((args :as list) (type-env :as dict) (prim-types :as dict) type-registry)
    ;; (if test then else?) → union of then and else types
    (if (< (len args) 2) "nil"
      (let ((then-type (infer-type (nth args 1) type-env prim-types type-registry)))
        (if (>= (len args) 3)
          (type-union then-type (infer-type (nth args 2) type-env prim-types type-registry))
          (type-union then-type "nil"))))))


(define infer-let-type
  (fn ((args :as list) (type-env :as dict) (prim-types :as dict) type-registry)
    ;; (let ((x expr) ...) body) → type of body in extended type-env
    (if (< (len args) 2) "nil"
      (let ((bindings (first args))
            (body (last args))
            (extended (merge type-env (dict))))
        ;; Add binding types
        (for-each
          (fn (binding)
            (when (and (= (type-of binding) "list") (>= (len binding) 2))
              (let ((name (if (= (type-of (first binding)) "symbol")
                            (symbol-name (first binding))
                            (str (first binding))))
                    (val-type (infer-type (nth binding 1) extended prim-types type-registry)))
                (dict-set! extended name val-type))))
          bindings)
        (infer-type body extended prim-types type-registry)))))


;; --------------------------------------------------------------------------
;; 7. Diagnostic types
;; --------------------------------------------------------------------------
;; Diagnostics are dicts:
;;   {:level "error"|"warning"|"info"
;;    :message "human-readable explanation"
;;    :component "~name" (or nil for top-level)
;;    :expr <the offending AST node>}

(define make-diagnostic
  (fn ((level :as string) (message :as string) component expr)
    {:level level
     :message message
     :component component
     :expr expr}))


;; --------------------------------------------------------------------------
;; 8. Call-site checking
;; --------------------------------------------------------------------------

(define check-primitive-call
  (fn ((name :as string) (args :as list) (type-env :as dict) (prim-types :as dict) prim-param-types (comp-name :as string) type-registry)
    ;; Check a primitive call site against declared param types.
    ;; prim-param-types is a dict: {prim-name → {:positional [...] :rest-type type-or-nil}}
    ;; Each positional entry is a list (name type-or-nil).
    ;; Returns list of diagnostics.
    (let ((diagnostics (list)))
      (when (and (not (nil? prim-param-types))
                 (has-key? prim-param-types name))
        (let ((sig (get prim-param-types name))
              (positional (get sig "positional"))
              (rest-type (get sig "rest-type")))
          ;; Check each positional arg
          (for-each
            (fn (idx)
              (when (< idx (len args))
                (if (< idx (len positional))
                  ;; Positional param — check against declared type
                  (let ((param-info (nth positional idx))
                        (arg-expr (nth args idx)))
                    (let ((expected-type (nth param-info 1)))
                      (when (not (nil? expected-type))
                        (let ((actual (infer-type arg-expr type-env prim-types type-registry)))
                          (when (and (not (type-any? expected-type))
                                     (not (type-any? actual))
                                     (not (subtype-resolved? actual expected-type type-registry)))
                            (append! diagnostics
                              (make-diagnostic "error"
                                (str "Argument " (+ idx 1) " of `" name
                                     "` expects " expected-type ", got " actual)
                                comp-name arg-expr)))))))
                  ;; Rest param — check against rest-type
                  (when (not (nil? rest-type))
                    (let ((arg-expr (nth args idx))
                          (actual (infer-type arg-expr type-env prim-types type-registry)))
                      (when (and (not (type-any? rest-type))
                                 (not (type-any? actual))
                                 (not (subtype-resolved? actual rest-type type-registry)))
                        (append! diagnostics
                          (make-diagnostic "error"
                            (str "Argument " (+ idx 1) " of `" name
                                 "` expects " rest-type ", got " actual)
                            comp-name arg-expr))))))))
            (range 0 (len args) 1))))
      diagnostics)))


(define check-component-call
  (fn ((comp-name :as string) (comp :as component) (call-args :as list) (type-env :as dict) (prim-types :as dict) type-registry)
    ;; Check a component call site against its declared param types.
    ;; comp is the component value, call-args is the list of args
    ;; from the call site (after the component name).
    (let ((diagnostics (list))
          (param-types (component-param-types comp))
          (params (component-params comp)))
      (when (and (not (nil? param-types))
                 (not (empty? (keys param-types))))
        ;; Parse keyword args from call site
        (let ((i 0)
              (provided-keys (list)))
          (for-each
            (fn (idx)
              (when (< idx (len call-args))
                (let ((arg (nth call-args idx)))
                  (when (= (type-of arg) "keyword")
                    (let ((key-name (keyword-name arg)))
                      (append! provided-keys key-name)
                      (when (< (+ idx 1) (len call-args))
                        (let ((val-expr (nth call-args (+ idx 1))))
                          ;; Check type of value against declared param type
                          (when (has-key? param-types key-name)
                            (let ((expected (get param-types key-name))
                                  (actual (infer-type val-expr type-env prim-types type-registry)))
                              (when (and (not (type-any? expected))
                                         (not (type-any? actual))
                                         (not (subtype-resolved? actual expected type-registry)))
                                (append! diagnostics
                                  (make-diagnostic "error"
                                    (str "Keyword :" key-name " of " comp-name
                                         " expects " expected ", got " actual)
                                    comp-name val-expr))))))))))))
            (range 0 (len call-args) 1))

          ;; Check for missing required params (those with declared types)
          (for-each
            (fn (param-name)
              (when (and (has-key? param-types param-name)
                         (not (contains? provided-keys param-name))
                         (not (type-nullable? (get param-types param-name))))
                (append! diagnostics
                  (make-diagnostic "warning"
                    (str "Required param :" param-name " of " comp-name " not provided")
                    comp-name nil))))
            params)

          ;; Check for unknown kwargs
          (for-each
            (fn (key)
              (when (not (contains? params key))
                (append! diagnostics
                  (make-diagnostic "warning"
                    (str "Unknown keyword :" key " passed to " comp-name)
                    comp-name nil))))
            provided-keys)))
      diagnostics)))


;; --------------------------------------------------------------------------
;; 9. AST walker — check a component body
;; --------------------------------------------------------------------------

(define check-body-walk
  (fn (node (comp-name :as string) (type-env :as dict) (prim-types :as dict) prim-param-types env (diagnostics :as list) type-registry effect-annotations)
    ;; Recursively walk an AST and collect diagnostics.
    ;; prim-param-types: dict of {name → {:positional [...] :rest-type t}} or nil
    ;; type-registry: dict of {type-name → type-def} or nil
    ;; effect-annotations: dict of {fn-name → effect-list} or nil
    (let ((kind (type-of node)))
      (when (= kind "list")
        (when (not (empty? node))
          (let ((head (first node))
                (args (rest node)))
            ;; Check calls when head is a symbol
            (when (= (type-of head) "symbol")
              (let ((name (symbol-name head)))
                ;; Component call
                (when (starts-with? name "~")
                  (let ((comp-val (env-get env name)))
                    (when (= (type-of comp-val) "component")
                      (for-each
                        (fn (d) (append! diagnostics d))
                        (check-component-call name comp-val args
                          type-env prim-types type-registry))))
                  ;; Effect check for component calls
                  (when (not (nil? effect-annotations))
                    (let ((caller-effects (get-effects comp-name effect-annotations)))
                      (for-each
                        (fn (d) (append! diagnostics d))
                        (check-effect-call name caller-effects effect-annotations comp-name)))))

                ;; Primitive call — check param types
                (when (and (not (starts-with? name "~"))
                           (not (nil? prim-param-types))
                           (has-key? prim-param-types name))
                  (for-each
                    (fn (d) (append! diagnostics d))
                    (check-primitive-call name args type-env prim-types
                      prim-param-types comp-name type-registry)))

                ;; Effect check for function calls
                (when (and (not (starts-with? name "~"))
                           (not (nil? effect-annotations)))
                  (let ((caller-effects (get-effects comp-name effect-annotations)))
                    (for-each
                      (fn (d) (append! diagnostics d))
                      (check-effect-call name caller-effects effect-annotations comp-name))))

                ;; Recurse into let with extended type env
                (when (or (= name "let") (= name "let*"))
                  (when (>= (len args) 2)
                    (let ((bindings (first args))
                          (body-exprs (rest args))
                          (extended (merge type-env (dict))))
                      (for-each
                        (fn (binding)
                          (when (and (= (type-of binding) "list")
                                     (>= (len binding) 2))
                            (let ((bname (if (= (type-of (first binding)) "symbol")
                                           (symbol-name (first binding))
                                           (str (first binding))))
                                  (val-type (infer-type (nth binding 1) extended prim-types type-registry)))
                              (dict-set! extended bname val-type))))
                        bindings)
                      (for-each
                        (fn (body)
                          (check-body-walk body comp-name extended prim-types prim-param-types env diagnostics type-registry effect-annotations))
                        body-exprs))))

                ;; Recurse into define with type binding
                (when (= name "define")
                  (when (>= (len args) 2)
                    (let ((def-name (if (= (type-of (first args)) "symbol")
                                     (symbol-name (first args))
                                     nil))
                          (def-val (nth args 1)))
                      (when def-name
                        (dict-set! type-env def-name
                          (infer-type def-val type-env prim-types type-registry)))
                      (check-body-walk def-val comp-name type-env prim-types prim-param-types env diagnostics type-registry effect-annotations))))))

            ;; Recurse into all child expressions
            (for-each
              (fn (child)
                (check-body-walk child comp-name type-env prim-types prim-param-types env diagnostics type-registry effect-annotations))
              args)))))))


;; --------------------------------------------------------------------------
;; 10. Check a single component
;; --------------------------------------------------------------------------

(define check-component
  (fn ((comp-name :as string) env (prim-types :as dict) prim-param-types type-registry effect-annotations)
    ;; Type-check a component's body. Returns list of diagnostics.
    ;; prim-param-types: dict of param type info, or nil to skip primitive checking.
    ;; type-registry: dict of {type-name → type-def} or nil
    ;; effect-annotations: dict of {fn-name → effect-list} or nil
    (let ((comp (env-get env comp-name))
          (diagnostics (list)))
      (when (= (type-of comp) "component")
        (let ((body (component-body comp))
              (params (component-params comp))
              (param-types (component-param-types comp))
              ;; Build initial type env from component params
              (type-env (dict)))
          ;; Add param types (annotated or default to any)
          (for-each
            (fn (p)
              (dict-set! type-env p
                (if (and (not (nil? param-types))
                         (has-key? param-types p))
                  (get param-types p)
                  "any")))
            params)
          ;; Add children as (list-of element) if component has children
          (when (component-has-children comp)
            (dict-set! type-env "children" (list "list-of" "element")))

          (check-body-walk body comp-name type-env prim-types prim-param-types env diagnostics type-registry effect-annotations)))
      diagnostics)))


;; --------------------------------------------------------------------------
;; 11. Check all components in an environment
;; --------------------------------------------------------------------------

(define check-all
  (fn (env (prim-types :as dict) prim-param-types type-registry effect-annotations)
    ;; Type-check every component in the environment.
    ;; prim-param-types: dict of param type info, or nil to skip primitive checking.
    ;; type-registry: dict of {type-name → type-def} or nil
    ;; effect-annotations: dict of {fn-name → effect-list} or nil
    ;; Returns list of all diagnostics.
    (let ((all-diagnostics (list)))
      (for-each
        (fn (name)
          (let ((val (env-get env name)))
            (when (= (type-of val) "component")
              (for-each
                (fn (d) (append! all-diagnostics d))
                (check-component name env prim-types prim-param-types type-registry effect-annotations)))))
        (keys env))
      all-diagnostics)))


;; --------------------------------------------------------------------------
;; 12. Build primitive type registry
;; --------------------------------------------------------------------------
;; Builds a dict mapping primitive-name → return-type from
;; the declarations parsed by boundary_parser.py.
;; This is called by the host at startup with the parsed declarations.

(define build-type-registry
  (fn ((prim-declarations :as list) (io-declarations :as list))
    ;; Both are lists of dicts: {:name "+" :returns "number" :params (...)}
    ;; Returns a flat dict: {"+" "number", "str" "string", ...}
    (let ((registry (dict)))
      (for-each
        (fn (decl)
          (let ((name (get decl "name"))
                (returns (get decl "returns")))
            (when (and (not (nil? name)) (not (nil? returns)))
              (dict-set! registry name returns))))
        prim-declarations)
      (for-each
        (fn (decl)
          (let ((name (get decl "name"))
                (returns (get decl "returns")))
            (when (and (not (nil? name)) (not (nil? returns)))
              (dict-set! registry name returns))))
        io-declarations)
      registry)))


;; --------------------------------------------------------------------------
;; 13. User-defined types (deftype)
;; --------------------------------------------------------------------------
;; Type definitions are plain dicts: {:name "price" :params [] :body "number"}
;; Stored in env under "*type-registry*" mapping type names to defs.

;; make-type-def and normalize-type-body are defined in eval.sx
;; (always compiled). They're available when types.sx is compiled as a spec module.

;; -- Standard type definitions --
;; These define the record types used throughout the type system itself.

;; Universal: nullable shorthand
(deftype (maybe a) (union a nil))

;; A type definition entry in the registry
(deftype type-def
  {:name string :params list :body any})

;; A diagnostic produced by the type checker
(deftype diagnostic
  {:level string :message string :component string? :expr any})

;; Primitive parameter type signature
(deftype prim-param-sig
  {:positional list :rest-type string?})

;; Effect declarations
(defeffect io)
(defeffect mutation)
(defeffect render)

(define type-def-name
  (fn (td) (get td "name")))

(define type-def-params
  (fn (td) (get td "params")))

(define type-def-body
  (fn (td) (get td "body")))

(define resolve-type
  (fn (t registry)
    ;; Resolve a type through the registry.
    ;; Returns the resolved type representation.
    (if (nil? registry) t
      (cond
        ;; String — might be a named type alias
        (= (type-of t) "string")
          (if (has-key? registry t)
            (let ((td (get registry t)))
              (let ((params (type-def-params td))
                    (body (type-def-body td)))
                (if (empty? params)
                  ;; Simple alias — resolve the body recursively
                  (resolve-type body registry)
                  ;; Parameterized with no args — return as-is
                  t)))
            t)
        ;; List — might be parameterized type application or compound
        (= (type-of t) "list")
          (if (empty? t) t
            (let ((head (first t)))
              (cond
                ;; (or ...), (list-of ...), (-> ...) — recurse into members
                (or (= head "or") (= head "list-of") (= head "->")
                    (= head "dict-of"))
                  (cons head (map (fn (m) (resolve-type m registry)) (rest t)))
                ;; Parameterized type application: ("maybe" "string") etc.
                (and (= (type-of head) "string")
                     (has-key? registry head))
                  (let ((td (get registry head))
                        (params (type-def-params td))
                        (body (type-def-body td))
                        (args (rest t)))
                    (if (= (len params) (len args))
                      (resolve-type
                        (substitute-type-vars body params args)
                        registry)
                      ;; Wrong arity — return as-is
                      t))
                :else t)))
        ;; Dict — record type, resolve field types
        (= (type-of t) "dict")
          (map-dict (fn (k v) (resolve-type v registry)) t)
        ;; Anything else — return as-is
        :else t))))

(define substitute-type-vars
  (fn (body (params :as list) (args :as list))
    ;; Substitute type variables in body.
    ;; params is a list of type var names, args is corresponding types.
    (let ((subst (dict)))
      (for-each
        (fn (i)
          (dict-set! subst (nth params i) (nth args i)))
        (range 0 (len params) 1))
      (substitute-in-type body subst))))

(define substitute-in-type
  (fn (t (subst :as dict))
    ;; Recursively substitute type variables.
    (cond
      (= (type-of t) "string")
        (if (has-key? subst t) (get subst t) t)
      (= (type-of t) "list")
        (map (fn (m) (substitute-in-type m subst)) t)
      (= (type-of t) "dict")
        (map-dict (fn (k v) (substitute-in-type v subst)) t)
      :else t)))

(define subtype-resolved?
  (fn (a b registry)
    ;; Resolve both sides through the registry, then check subtype.
    (if (nil? registry)
      (subtype? a b)
      (let ((ra (resolve-type a registry))
            (rb (resolve-type b registry)))
        ;; Handle record structural subtyping: dict a <: dict b
        ;; if every field in b exists in a with compatible type
        (if (and (= (type-of ra) "dict") (= (type-of rb) "dict"))
          (every?
            (fn (key)
              (and (has-key? ra key)
                   (subtype-resolved? (get ra key) (get rb key) registry)))
            (keys rb))
          (subtype? ra rb))))))


;; --------------------------------------------------------------------------
;; 14. Effect checking (defeffect)
;; --------------------------------------------------------------------------
;; Effects are annotations on functions/components describing their
;; side effects. A pure function cannot call IO functions.

(define get-effects
  (fn ((name :as string) effect-annotations)
    ;; Look up declared effects for a function/component.
    ;; Returns list of effect strings, or nil if unannotated.
    (if (nil? effect-annotations) nil
      (if (has-key? effect-annotations name)
        (get effect-annotations name)
        nil))))

(define effects-subset?
  (fn (callee-effects caller-effects)
    ;; Are all callee effects allowed by caller?
    ;; nil effects = unannotated = assumed to have all effects.
    ;; Empty list = pure = no effects.
    (if (nil? caller-effects) true     ;; unannotated caller allows everything
      (if (nil? callee-effects) true   ;; unannotated callee — skip check
        (every?
          (fn (e) (contains? caller-effects e))
          callee-effects)))))

(define check-effect-call
  (fn ((callee-name :as string) caller-effects effect-annotations (comp-name :as string))
    ;; Check that callee's effects are allowed by caller's effects.
    ;; Returns list of diagnostics.
    (let ((diagnostics (list))
          (callee-effects (get-effects callee-name effect-annotations)))
      (when (and (not (nil? caller-effects))
                 (not (nil? callee-effects))
                 (not (effects-subset? callee-effects caller-effects)))
        (append! diagnostics
          (make-diagnostic "error"
            (str "`" callee-name "` has effects "
                 (join ", " callee-effects)
                 " but `" comp-name "` only allows "
                 (if (empty? caller-effects) "[pure]"
                   (join ", " caller-effects)))
            comp-name nil)))
      diagnostics)))

(define build-effect-annotations
  (fn ((io-declarations :as list))
    ;; Assign [io] effect to all IO primitives.
    (let ((annotations (dict)))
      (for-each
        (fn (decl)
          (let ((name (get decl "name")))
            (when (not (nil? name))
              (dict-set! annotations name (list "io")))))
        io-declarations)
      annotations)))


;; --------------------------------------------------------------------------
;; 15. Check component effects — convenience wrapper
;; --------------------------------------------------------------------------
;; Validates that components respect their declared effect annotations.
;; Delegates to check-body-walk with nil type checking (effects only).

(define check-component-effects
  (fn ((comp-name :as string) env effect-annotations)
    ;; Check a single component's effect usage. Returns diagnostics list.
    ;; Skips type checking — only checks effect violations.
    (let ((comp (env-get env comp-name))
          (diagnostics (list)))
      (when (= (type-of comp) "component")
        (let ((body (component-body comp)))
          (check-body-walk body comp-name (dict) (dict) nil env
            diagnostics nil effect-annotations)))
      diagnostics)))

(define check-all-effects
  (fn (env effect-annotations)
    ;; Check all components in env for effect violations.
    ;; Returns list of all diagnostics.
    (let ((all-diagnostics (list)))
      (for-each
        (fn (name)
          (let ((val (env-get env name)))
            (when (= (type-of val) "component")
              (for-each
                (fn (d) (append! all-diagnostics d))
                (check-component-effects name env effect-annotations)))))
        (keys env))
      all-diagnostics)))


;; --------------------------------------------------------------------------
;; Platform interface summary
;; --------------------------------------------------------------------------
;;
;; From eval.sx (already provided):
;;   (type-of x), (symbol-name s), (keyword-name k), (env-get env k)
;;   (component-body c), (component-params c), (component-has-children c)
;;
;; New for types.sx (each host implements):
;;   (component-param-types c)        → dict {param-name → type} or nil
;;   (component-set-param-types! c d) → store param types on component
;;   (merge d1 d2)               → new dict merging d1 and d2
;;
;; Primitive param types:
;;   The host provides prim-param-types as a dict mapping primitive names
;;   to param type descriptors. Each descriptor is a dict:
;;     {"positional" [["name" "type-or-nil"] ...] "rest-type" "type-or-nil"}
;;   Built by boundary_parser.parse_primitive_param_types() in Python.
;;   Passed to check-component/check-all as an optional extra argument.
;;
;; --------------------------------------------------------------------------