#lang racket/base

(provide (struct-out ty-union)
         (struct-out ty-unit)
         (struct-out ty-value)
         (struct-out ty-record)
         (struct-out ty-array)
         (struct-out ty-set)
         (struct-out ty-dictionary)

         definition-ty
         add-name-if-absent
         escape)

(require preserves/record)
(require racket/match)
(require (only-in racket/syntax format-symbol))

(struct ty-union (variants) #:transparent)
(struct ty-unit () #:transparent)
(struct ty-value () #:transparent)
(struct ty-record (fields) #:transparent)
(struct ty-array (type) #:transparent)
(struct ty-set (type) #:transparent)
(struct ty-dictionary (key-type value-type) #:transparent)

(define (definition-ty d)
  (match d
    [(record 'or (list named-alts))
     (ty-union (map (match-lambda
                      [(list variant-label-str pattern)
                       (list (string->symbol variant-label-str) (pattern-ty pattern))])
                    named-alts))]
    [(record 'and (list named-pats)) (product-ty named-pats)]
    [pattern (pattern-ty pattern)]))

(define (product-ty named-pats)
  (match (gather-fields* named-pats '())
    ['() (ty-unit)]
    [fields (ty-record fields)]))

(define (gather-fields* named-pats acc)
  (foldr gather-fields acc named-pats))

(define (gather-fields named-pat acc)
  (match named-pat
    [(record 'named (list n p))
     (match (pattern-ty p)
       [(ty-unit) acc]
       [ty (cons (list n ty) acc)])]
    [(record 'rec (list label-named-pat fields-named-pat))
     (gather-fields label-named-pat (gather-fields fields-named-pat acc))]
    [(record 'tuple (list named-pats)) (gather-fields* named-pats acc)]
    [(record 'tuple* (list fixed-named-pats variable-named-pat))
     (gather-fields* fixed-named-pats (gather-fields variable-named-pat acc))]
    [(record 'dict (list (hash-table (keys pats) ...)))
     (gather-fields* (map add-name-if-absent keys pats) acc)]
    [_ acc]))

(define (pattern-ty p)
  (match p
    ['any (ty-value)]
    [(record 'atom (list _atom-kind)) (ty-value)]
    [(record 'embedded '()) (ty-value)]
    [(record 'lit (list _value)) (ty-unit)]
    [(record 'ref (list _module-path _name)) (ty-value)]
    [(record 'tuple* (list '() (and variable-pat (not (record 'named _)))))
     (ty-array (pattern-ty variable-pat))]
    [(record 'setof (list pat)) (ty-set (pattern-ty pat))]
    [(record 'dictof (list key-pat value-pat))
     (ty-dictionary (pattern-ty key-pat) (pattern-ty value-pat))]
    [_ (product-ty (list p))]))

(define (add-name-if-absent k p)
  (match p
    [(record 'named _) p]
    [_ (match (namelike k)
         [#f p]
         [s (record 'named (list s p))])]))

(define (namelike v)
  (match v
    [(? string? s) (string->symbol s)]
    [(? symbol? s) s]
    [(? number? n) (string->symbol (number->string n))]
    [(? boolean? b) (if b 'true 'false)]
    [_ #f]))

(define (escape s)
  (format-symbol "$~a" s))