#lang scheme
(require (for-syntax syntax/parse)
         scheme/generator)

(define-struct model (rules))
(define-struct rule (head body))
(define-struct state (env qs))
(define-struct query ())
(define-struct (sexpr-query query) (se))
(define-struct (fun-query query) (f args))

(define (variable? q)
  (and (symbol? q)
       (char-upper-case?
        (string-ref 
         (symbol->string q)
         0))))

(define (unbound-variable? env q1)
  (and (variable? q1)
       (not (hash-has-key? env q1))))

(define (bound-variable? env q1)
  (and (variable? q1)
       (hash-has-key? env q1)))

(define (unify env q1 q2)
  (cond
    [(equal? q1 q2)
     env]
    [(unbound-variable? env q1)
     (hash-set env q1 q2)]
    [(unbound-variable? env q2)
     (hash-set env q2 q1)]
    [(bound-variable? env q1)
     (unify env (hash-ref env q1) q2)]
    [(bound-variable? env q2)
     (unify env q1 (hash-ref env q2))]
    [(and (list? q1) (list? q2)
          (= (length q1) (length q2)))
     (for/fold ([env env])
       ([e1 (in-list q1)]
        [e2 (in-list q2)])
       (and env
            (unify env e1 e2)))]
    [else #f]))

(define alpha-vary
  (match-lambda
    [(struct rule (head body))
     (define mapping (make-hasheq))
     (define alpha-se
       (match-lambda
         [(? variable? v)
          (hash-ref! mapping v (lambda () (gensym v)))]
         [(list-rest e1 e2)
          (cons (alpha-se e1)
                (alpha-se e2))]
         [e
          e]))
     (define alpha
       (match-lambda
         [(struct sexpr-query (q))
          (make-sexpr-query (alpha-se q))]
         [(struct fun-query (f args))
          (make-fun-query f (alpha-se args))]))
     (make-rule (alpha-se head)
                (map alpha body))]))

(define (rule-matches? r env q)
  (match-define (struct rule (head body)) r)
  (define new-env (unify env head q))
  (and new-env
       (make-state new-env body)))

(define generator-done
  (local [(define-struct uniq ())]
    (make-uniq)))

(define (model-state-generator m s)
  (match-define (struct state (env (list-rest first-query other-queries))) s)
  (generator
   (match first-query
     [(struct sexpr-query (q-se))
      (for ([rule (in-list (model-rules m))])
        (match (rule-matches? (alpha-vary rule) env q-se)
          [#f (void)]
          [(struct state (new-env subqueries))
           (yield (make-state new-env (append subqueries other-queries)))]))]
     [(struct fun-query (f args))
      (when (apply f (map (curry env-deref env) args))
        (yield (make-state env other-queries)))])
   (yield generator-done)))

(define (model-final-state-generator m t)
  (generator
   (match t
     [(struct state (env (list)))
      (yield env)]
     [(? state?)
      (for* ([next-state (in-producer (model-state-generator m t) generator-done)]
             [ans (in-producer (model-final-state-generator m next-state) generator-done)])
        (yield ans))])
   (yield generator-done)))

(define (env-deref env v)
  (cond
    [(bound-variable? env v)
     (env-deref env (hash-ref env v))]
    [(list? v)
     (map (curry env-deref env) v)]
    [else
     v]))

(define (env-restrict env l)
  (for/hasheq ([(k v) (in-hash env)]
               #:when (member k l))
    (values k (env-deref env v))))

(define (variables-in q)
  (match q
    [(? variable? v) (list v)]
    [(struct fun-query (_ l))
     (append-map variables-in l)]
    [(struct sexpr-query (l))
     (append-map variables-in l)]
    [_ empty]))

(define (query-answer-generator m q)
  (define init-vars (variables-in q))
  (generator
   (for ([ans (in-producer (model-final-state-generator 
                            m 
                            (make-state (make-immutable-hasheq empty)
                                        (list q)))
                           generator-done)])
     (yield (env-restrict ans init-vars)))
   (yield generator-done)))

(define (query-model* m q #:limit [limit +inf.0])
  (for/list ([ans (in-producer (query-answer-generator m q) generator-done)]
             [i (in-range limit)])
    ans))

(define-syntax (compile-query stx)
  (syntax-parse
   stx #:literals (unquote)
   [(_ ((unquote f) arg ...))
    (syntax/loc stx
      (make-fun-query f (list 'arg ...)))]
   [(_ query)
    (syntax/loc stx
      (make-sexpr-query 'query))]))

(provide (all-defined-out))