#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
         unwrap
         add-name-if-absent
         escape)

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

(require "gen/schema.rkt")

(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
    [(Definition-or p0 p1 pN)
     (ty-union (map (match-lambda
                      [(NamedAlternative variant-label-str pattern)
                       (list (string->symbol variant-label-str) (pattern-ty pattern))])
                    (list* p0 p1 pN)))]
    [(Definition-and p0 p1 pN) (product-ty (list* p0 p1 pN))]
    [(Definition-Pattern pattern) (pattern-ty pattern)]))

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

(define (unwrap p)
  (match p
    [(Pattern-SimplePattern p) (unwrap p)]
    [(Pattern-CompoundPattern p) (unwrap p)]
    [(NamedPattern-named p) (unwrap p)]
    [(NamedSimplePattern-named p) (unwrap p)]
    [(NamedPattern-anonymous p) (unwrap p)]
    [(NamedSimplePattern-anonymous p) (unwrap p)]
    [_ p]))

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

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

(define (pattern-ty p)
  (match (unwrap p)
    [(SimplePattern-any) (ty-value)]
    [(SimplePattern-atom _atomKind) (ty-value)]
    [(SimplePattern-embedded) (ty-value)]
    [(SimplePattern-lit _value) (ty-unit)]
    [(SimplePattern-Ref _r) (ty-value)]
    [(CompoundPattern-tuple* '() (NamedSimplePattern-anonymous variable-pat))
     (ty-array (pattern-ty variable-pat))]
    [(CompoundPattern-setof pat) (ty-set (pattern-ty pat))]
    [(CompoundPattern-dictof key-pat value-pat)
     (ty-dictionary (pattern-ty key-pat) (pattern-ty value-pat))]
    [(? CompoundPattern?) (product-ty (list p))]))

(define (add-name-if-absent k p)
  (match p
    [(NamedSimplePattern-named _) p]
    [(NamedSimplePattern-anonymous _)
     (match (namelike k)
       [#f p]
       [s (NamedSimplePattern-named (NamedSimplePattern_ 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))