-- Inspired by https://github.com/coq-community/reglang/blob/master/theories/nfa.v

import Clash.Prelude
import qualified Prelude
import qualified Data.List

data Re a
  = Eps
  | Char (a -> Bool)
  | Cat (Re a) (Re a)
  | Alt (Re a) (Re a)
  | Star (Re a)

-- ============================================================

lshift :: [a] -> [Either a b]
lshift = Prelude.map Left

rshift :: [b] -> [Either a b]
rshift = Prelude.map Right

-------------------------------------------------------------

data Nfa a = Nfa {
    start :: [a]
  , final :: [a]
  , tfun :: a -> Maybe Char -> a -> Bool
}

-------------------------------------------------------------

eps :: Nfa ()
eps = Nfa {
    start = [()]
  , final = [()]
  , tfun = \ s c d -> case (c, s, d) of
      (Nothing, (), ()) -> True
      otherwise -> False
    -- () to () possible but not by consuming `c'
}

char :: (Char -> Bool) -> Nfa Bool
char f = Nfa {
    start = [False]
  , final = [True]
  , tfun = \ s c d -> case (c, s, d) of
      (Just c, False, True) -> f c
      (Nothing, False, False) -> True
      (Nothing, True, True) -> True
      otherwise -> False
}

cat :: (Eq a, Eq b) => Nfa a -> Nfa b -> Nfa (Either a b)
cat a1 a2 = Nfa {
    start = lshift (start a1)
  , final = rshift (final a2)
  , tfun = \ s c d -> case (c, s, d) of
      (Just c', Left s, Left d) -> (tfun a1) s c d
      (Just c', Right s, Right d) -> (tfun a2) s c d

      -- The epsilon transition
      (Nothing, Left s, Right d) ->
        if Prelude.elem s (final a1) && Prelude.elem d (start a2) then True
        else False
      otherwise -> False
}

-- if ∃st ∈ (start a1) and st ∈ (final a1) then
--   (start a1) ++ 
-- else
--   start a1

alt :: Nfa a -> Nfa b -> Nfa (Either a b)
alt a1 a2 = Nfa {
    start = (Data.List.++) (lshift (start a1)) (rshift (start a2))
  , final = (Data.List.++) (lshift (final a1)) (rshift (final a2))
  , tfun = \ s c d -> case (c, s, d) of
      (_, Left s, Left d) -> (tfun a1) s c d
      (_, Right s, Right d) -> (tfun a2) s c d
      otherwise -> False
}

star :: Eq a => Nfa a -> Nfa a
star au = Nfa {
    start = start au
  , final = (Data.List.++) (start au) (final au)
  , tfun = \ s c d -> case (c, s, d) of
      (Nothing, _, _) ->
        if Prelude.elem s (final au) && Prelude.elem d (start au) then True
        else False
      otherwise -> (tfun au) s c d
}

-------------------------------------------------------------

-- re2nfa :: Re a -> Nfa st
-- re2nfa r = case r of
--   Eps -> eps
--   Char f -> char f
--   Cat r1 r2 -> cat (re2nfa r1) (re2nfa r2)
--   Alt r1 r2 -> alt (re2nfa r1) (re2nfa r2)
--   Star rr -> star (re2nfa rr)