(in-package #:nekod.regex)

;;; NFA and DFA compilation from regex algebra
;;; Pipeline: REGEX_AST -> NFA -> epsilon-closure -> subset construction -> DFA

;;; --- NFA representation ---

(defstruct nfa-state
  (id 0 :type integer)
  (transitions nil :type list)
  (epsilon nil :type list)
  (accepting nil :type boolean))

(defstruct nfa
  (states (make-array 0 :adjustable t :fill-pointer 0))
  (initial 0 :type integer)
  (accepting 0 :type integer))

(defvar *nfa-state-counter* 0)

(defun fresh-nfa-state (&optional accepting)
  (let ((s (make-nfa-state :id (incf *nfa-state-counter*)
                           :accepting accepting)))
    s))

;;; --- Thompson's construction: regex -> NFA ---

(defun regex-to-nfa (re)
  "Convert regex algebra tree to NFA via Thompson's construction."
  (let ((*nfa-state-counter* 0))
    (regex-to-nfa-inner re)))

(defun regex-to-nfa-inner (re)
  (etypecase re
    (re-empty
     (let ((s (fresh-nfa-state nil))
           (f (fresh-nfa-state t)))
       (make-nfa :states (make-array 2 :initial-contents (list s f) :adjustable t :fill-pointer 2)
                 :initial (nfa-state-id s)
                 :accepting (nfa-state-id f))))
    (re-epsilon
     (let ((s (fresh-nfa-state nil))
           (f (fresh-nfa-state t)))
       (push (nfa-state-id f) (nfa-state-epsilon s))
       (make-nfa :states (make-array 2 :initial-contents (list s f) :adjustable t :fill-pointer 2)
                 :initial (nfa-state-id s)
                 :accepting (nfa-state-id f))))
    (re-literal
     (let* ((chars (re-literal-value re))
            (states (loop for i from 0 to (length chars)
                          collect (fresh-nfa-state (= i (length chars))))))
       (loop for i from 0 below (length chars)
             for from-state = (nth i states)
             for to-state = (nth (1+ i) states)
             do (push (cons (char chars i) (nfa-state-id to-state))
                      (nfa-state-transitions from-state)))
       (make-nfa :states (make-array (length states) :initial-contents states
                                     :adjustable t :fill-pointer (length states))
                 :initial (nfa-state-id (first states))
                 :accepting (nfa-state-id (car (last states))))))
    (re-concat
     (let ((left (regex-to-nfa-inner (re-concat-left re)))
           (right (regex-to-nfa-inner (re-concat-right re))))
       (let ((left-accept (find (nfa-accepting left) (nfa-states left)
                                :key #'nfa-state-id)))
         (setf (nfa-state-accepting left-accept) nil)
         (push (nfa-initial right) (nfa-state-epsilon left-accept)))
       (let ((all-states (concatenate 'vector (nfa-states left) (nfa-states right))))
         (make-nfa :states (make-array (length all-states) :initial-contents all-states
                                       :adjustable t :fill-pointer (length all-states))
                   :initial (nfa-initial left)
                   :accepting (nfa-accepting right)))))
    (re-union
     (let ((left (regex-to-nfa-inner (re-union-left re)))
           (right (regex-to-nfa-inner (re-union-right re)))
           (start (fresh-nfa-state nil))
           (final (fresh-nfa-state t)))
       (push (nfa-initial left) (nfa-state-epsilon start))
       (push (nfa-initial right) (nfa-state-epsilon start))
       (let ((left-accept (find (nfa-accepting left) (nfa-states left)
                                :key #'nfa-state-id))
             (right-accept (find (nfa-accepting right) (nfa-states right)
                                 :key #'nfa-state-id)))
         (setf (nfa-state-accepting left-accept) nil)
         (setf (nfa-state-accepting right-accept) nil)
         (push (nfa-state-id final) (nfa-state-epsilon left-accept))
         (push (nfa-state-id final) (nfa-state-epsilon right-accept)))
       (let ((all (concatenate 'vector
                               (vector start final)
                               (nfa-states left)
                               (nfa-states right))))
         (make-nfa :states (make-array (length all) :initial-contents all
                                       :adjustable t :fill-pointer (length all))
                   :initial (nfa-state-id start)
                   :accepting (nfa-state-id final)))))
    (re-kleene
     (let ((inner (regex-to-nfa-inner (re-kleene-inner re)))
           (start (fresh-nfa-state nil))
           (final (fresh-nfa-state t)))
       (push (nfa-initial inner) (nfa-state-epsilon start))
       (push (nfa-state-id final) (nfa-state-epsilon start))
       (let ((inner-accept (find (nfa-accepting inner) (nfa-states inner)
                                 :key #'nfa-state-id)))
         (setf (nfa-state-accepting inner-accept) nil)
         (push (nfa-initial inner) (nfa-state-epsilon inner-accept))
         (push (nfa-state-id final) (nfa-state-epsilon inner-accept)))
       (let ((all (concatenate 'vector
                               (vector start final)
                               (nfa-states inner))))
         (make-nfa :states (make-array (length all) :initial-contents all
                                       :adjustable t :fill-pointer (length all))
                   :initial (nfa-state-id start)
                   :accepting (nfa-state-id final)))))
    (re-plus
     (regex-to-nfa-inner (re-concat (re-plus-inner re)
                                    (re-kleene (re-plus-inner re)))))
    (re-optional
     (regex-to-nfa-inner (re-union re (make-re-epsilon))))))

;;; --- Epsilon closure ---

(defun find-nfa-state (nfa id)
  (find id (nfa-states nfa) :key #'nfa-state-id))

(defun epsilon-closure (nfa state-ids)
  "Compute epsilon-closure of a set of NFA state IDs."
  (let ((closure (make-hash-table))
        (stack (copy-list state-ids)))
    (dolist (s state-ids) (setf (gethash s closure) t))
    (loop while stack
          do (let* ((sid (pop stack))
                    (state (find-nfa-state nfa sid)))
               (when state
                 (dolist (eps (nfa-state-epsilon state))
                   (unless (gethash eps closure)
                     (setf (gethash eps closure) t)
                     (push eps stack))))))
    (let ((result nil))
      (maphash (lambda (k v) (declare (ignore v)) (push k result)) closure)
      (sort result #'<))))

(defun nfa-move (nfa state-ids char)
  "Compute set of states reachable from state-ids on input char."
  (let ((result nil))
    (dolist (sid state-ids)
      (let ((state (find-nfa-state nfa sid)))
        (when state
          (dolist (tr (nfa-state-transitions state))
            (when (char= (car tr) char)
              (pushnew (cdr tr) result))))))
    result))

;;; --- Subset construction: NFA -> DFA ---

(defstruct (dfa-state (:constructor make-dfa-state-struct))
  id
  transitions
  accepting)

(defstruct dfa
  states
  initial
  accepting)

(defun make-dfa-state (id &optional accepting)
  (make-dfa-state-struct :id id
                         :transitions (make-hash-table :test #'equal)
                         :accepting accepting))

(defun dfa-add-transition (dfa from-id char to-id)
  (let ((from-state (aref (dfa-states dfa) from-id)))
    (setf (gethash char (dfa-state-transitions from-state)) to-id)))

(defun dfa-step (dfa state-id input-char)
  (let ((state (aref (dfa-states dfa) state-id)))
    (gethash input-char (dfa-state-transitions state))))

(defun dfa-run (dfa input)
  "Run DFA on input string. Returns T if accepted."
  (let ((current-state (dfa-initial dfa)))
    (loop for c across input
          do (let ((next (dfa-step dfa current-state c)))
               (if next
                   (setf current-state next)
                   (return-from dfa-run nil))))
    (dfa-state-accepting (aref (dfa-states dfa) current-state))))

(defun nfa-to-dfa (nfa)
  "Convert NFA to DFA via subset construction."
  (let* ((initial-closure (epsilon-closure nfa (list (nfa-initial nfa))))
         (dfa-states-list nil)
         (state-map (make-hash-table :test #'equal))
         (state-counter 0)
         (worklist nil)
         (alphabet (collect-alphabet nfa)))
    (setf (gethash initial-closure state-map) state-counter)
    (push (cons initial-closure state-counter) dfa-states-list)
    (incf state-counter)
    (push initial-closure worklist)
    (let ((transitions nil))
      (loop while worklist
            do (let* ((current (pop worklist))
                      (current-id (gethash current state-map)))
                 (dolist (c alphabet)
                   (let* ((moved (nfa-move nfa current c))
                          (target (and moved (epsilon-closure nfa moved))))
                     (when target
                       (unless (gethash target state-map)
                         (setf (gethash target state-map) state-counter)
                         (push (cons target state-counter) dfa-states-list)
                         (incf state-counter)
                         (push target worklist))
                       (push (list current-id c (gethash target state-map))
                             transitions))))))
      (let* ((n state-counter)
             (states (make-array n)))
        (loop for i from 0 below n
              do (let* ((entry (find i dfa-states-list :key #'cdr))
                        (nfa-ids (car entry))
                        (accepting (some (lambda (sid)
                                          (let ((s (find-nfa-state nfa sid)))
                                            (and s (nfa-state-accepting s))))
                                        nfa-ids)))
                   (setf (aref states i) (make-dfa-state i accepting))))
        (let ((result (make-dfa :states states :initial 0 :accepting nil)))
          (dolist (tr transitions)
            (dfa-add-transition result (first tr) (second tr) (third tr)))
          (setf (dfa-accepting result)
                (loop for i from 0 below n
                      when (dfa-state-accepting (aref states i))
                      collect i))
          result)))))

(defun collect-alphabet (nfa)
  "Collect all characters used in NFA transitions."
  (let ((chars nil))
    (loop for state across (nfa-states nfa)
          do (dolist (tr (nfa-state-transitions state))
               (pushnew (car tr) chars)))
    chars))

;;; --- Top-level compilation ---

(defun compile-regex-to-dfa (re)
  "Compile regex algebra tree directly to DFA. No runtime regex needed."
  (nfa-to-dfa (regex-to-nfa re)))