#lang racket/base

(provide (for-syntax preserves-pattern-registry
                     register-preserves-pattern!
                     analyse-pattern
                     analyse-pattern-bindings)
         define-preserves-pattern
         :pattern
         :template)

(require (for-syntax racket/base))
(require (for-syntax racket/match))
(require (for-syntax racket/list))
(require (for-syntax racket/struct-info))
(require (for-syntax racket/syntax))
(require (for-syntax syntax/id-table))
(require (for-syntax syntax/stx))

(require "pattern-expander.rkt")
(require "schemas/gen/dataspace-patterns.rkt")

;;---------------------------------------------------------------------------

(define-for-syntax preserves-pattern-registry (make-free-id-table))

(define-for-syntax (register-preserves-pattern! id-stx transformer)
  (free-id-table-set! preserves-pattern-registry id-stx transformer))

(define-syntax (define-preserves-pattern stx)
  (syntax-case stx ()
    [(_ (ctor-stx field-stxs ...) pattern-stx bindings-stx)
     #`(begin (begin-for-syntax
                (register-preserves-pattern!
                 #'ctor-stx
                 (lambda (mode s)
                   (match mode
                     ['pattern
                      (syntax-case s ()
                        [(_ field-stxs ...) pattern-stx]
                        [_ (raise-syntax-error 'ctor-stx "Invalid pattern")])]
                     ['template
                      (syntax-case s ()
                        [(_ field-stxs ...)
                         (syntax (#,(format-id #'ctor-stx "~a->preserves" (syntax-e #'ctor-stx))
                                  (ctor-stx field-stxs ...)))])]
                     ['bindings
                      (syntax-case s ()
                        [(_ field-stxs ...) bindings-stx]
                        [_ (raise-syntax-error 'ctor-stx "Invalid binding-pattern")])]))))
              (void))]))

;;---------------------------------------------------------------------------

(define-for-syntax orig-insp
  (variable-reference->module-declaration-inspector (#%variable-reference)))

(begin-for-syntax
  (define (dollar-id? stx)
    (and (identifier? stx)
         (char=? (string-ref (symbol->string (syntax-e stx)) 0) #\$)))

  (define (undollar stx)
    (and (dollar-id? stx)
         (datum->syntax stx (string->symbol (substring (symbol->string (syntax-e stx)) 1)))))

  (define (discard-id? stx)
    (and (identifier? stx)
         (free-identifier=? #'_ stx)))

  (define (id-value stx)
    (and (identifier? stx)
         (syntax-local-value stx (lambda () #f))))

  (define (list-id? stx)
    (and (identifier? stx)
         (free-identifier=? #'list stx)))

  (define (constructor-registered? stx)
    (free-id-table-ref preserves-pattern-registry stx #f))

  (define (member-entry key-stx pat-stx)
    (define analysed (analyse-pattern pat-stx))
    (syntax-case analysed (DDiscard)
      [(DDiscard) (list)]
      [_ (list key-stx analysed)]))

  (define (struct-label-and-arity id-stx actual-count)
    (match-define (list _type-stx ctor-stx pred-stx accessor-stxs _mutator-stxs _super)
      (extract-struct-info (id-value id-stx)))
    (define expected-count (length accessor-stxs))
    (when (not (= expected-count actual-count))
      (raise-syntax-error ':pattern
                          (format "Expected ~a arguments to ~v, but got ~a"
                                  expected-count
                                  ctor-stx
                                  actual-count)))
    (values (syntax-e ctor-stx)
            expected-count))

  (define (analyse-pattern stx)
    (define disarmed-stx (syntax-disarm stx orig-insp))
    (syntax-case disarmed-stx ($ quasiquote unquote quote)
      [(ctor args ...)
       (constructor-registered? #'ctor)
       ((free-id-table-ref preserves-pattern-registry #'ctor) 'pattern disarmed-stx)]

      [(expander args ...)
       (pattern-expander-id? #'expander)
       (pattern-expander-transform disarmed-stx
                                   (lambda (result)
                                     (analyse-pattern (syntax-rearm result stx))))]

      ;; Extremely limited support for quasiquoting and quoting
      [(quasiquote (unquote p)) (analyse-pattern #'p)]
      [(quasiquote (p ...)) (analyse-pattern #'(list (quasiquote p) ...))]
      [(quasiquote p) #`(DLit 'p)]
      [(quote p) #`(DLit 'p)]

      [(ctor piece ...)
       (struct-info? (id-value #'ctor))
       (let-values (((label arity) (struct-label-and-arity #'ctor)))
         #`(DCompound-rec '#,label
                          #,arity
                          (hasheqv #,@(append*
                                       (for/list [(n (in-naturals))
                                                  (piece (in-list (syntax->list #'(piece ...))))]
                                         (member-entry n piece))))))]

      [(list-stx piece ...)
       (list-id? #'list-stx)
       #`(DCompound-arr #,(length (syntax->list #'(piece ...)))
                        (hasheqv #,@(append*
                                     (for/list [(n (in-naturals))
                                                (piece (in-list (syntax->list #'(piece ...))))]
                                       (member-entry n piece)))))]

      [id
       (dollar-id? #'id)
       #`(DBind '#,(undollar #'id) (DDiscard))]

      [($ id p)
       #`(DBind 'id #,(analyse-pattern #'p))]

      [id
       (discard-id? #'id)
       #`(DDiscard)]

      [other
       #`(DLit other)]))

  (define (analyse-pattern-bindings stx)
    (define disarmed-stx (syntax-disarm stx orig-insp))
    (syntax-case disarmed-stx ($ quasiquote unquote quote)
      [(ctor args ...)
       (constructor-registered? #'ctor)
       ((free-id-table-ref preserves-pattern-registry #'ctor) 'bindings disarmed-stx)]

      [(expander args ...)
       (pattern-expander-id? #'expander)
       (pattern-expander-transform disarmed-stx
                                   (lambda (result)
                                     (analyse-pattern-bindings (syntax-rearm result stx))))]

      ;; Extremely limited support for quasiquoting and quoting
      [(quasiquote (unquote p)) (analyse-pattern-bindings #'p)]
      [(quasiquote (p ...)) (analyse-pattern-bindings #'(list (quasiquote p) ...))]
      [(quasiquote _p) '()]
      [(quote _p) '()]

      [(ctor piece ...)
       (struct-info? (id-value #'ctor))
       (append-map analyse-pattern-bindings (syntax->list #'(piece ...)))]

      [(list-stx piece ...)
       (list-id? #'list-stx)
       (append-map analyse-pattern-bindings (syntax->list #'(piece ...)))]

      [id
       (dollar-id? #'id)
       (list (undollar #'id))]

      [($ id p)
       (cons #'id (analyse-pattern-bindings #'p))]

      [id
       (discard-id? #'id)
       '()]

      [other
       '()]))

  (define (analyse-template stx)
    (syntax-case stx ($ quasiquote unquote quote)
      [(ctor args ...)
       (constructor-registered? #'ctor)
       ((free-id-table-ref preserves-pattern-registry #'ctor) 'template stx)]

      ;; Extremely limited support for quasiquoting and quoting
      [(quasiquote (unquote p)) #'p]
      [(quasiquote (p ...)) (analyse-template #'(list (quasiquote p) ...))]
      [(quasiquote p) #''p]
      [(quote p) #''p]

      [(ctor piece ...)
       (struct-info? (id-value #'ctor))
       #`(ctor #,@(map analyse-template (syntax->list #'(piece ...))))]

      [(list-stx piece ...)
       (list-id? #'list-stx)
       #`(list-stx #,@(map analyse-template (syntax->list #'(piece ...))))]

      [other #'other])))

(define-syntax (:pattern stx)
  (syntax-case stx ()
    [(_ pat-stx)
     (analyse-pattern #'pat-stx)]))

;; (provide :bindings)
;; (define-syntax (:bindings stx)
;;   (syntax-case stx ()
;;     [(_ pat-stx)
;;      #`(quote #,(analyse-pattern-bindings #'pat-stx))]))

(define-syntax (:template stx)
  (syntax-case stx ()
    [(_ template-stx)
     (analyse-template #'template-stx)]))