{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE StandaloneDeriving #-} -- For Show (ParsingError inp)
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UnboxedTuples #-} -- For nextInput
{-# LANGUAGE UndecidableInstances #-} -- For Show (ParsingError inp)
module Symantic.Parser.Machine.Generate where

import Control.Monad (Monad(..))
import Data.Bool (Bool)
import Data.Char (Char)
import Data.Either (Either(..))
import Data.Function (($))
-- import Data.Functor ((<$>))
import Data.Int (Int)
import Data.Maybe (Maybe(..))
import Data.Ord (Ord, Ordering(..))
import Data.Semigroup (Semigroup(..))
import Data.Set (Set)
import Language.Haskell.TH (CodeQ, Code(..))
import Prelude (($!))
import Text.Show (Show(..))
import qualified Data.Eq as Eq
import qualified Data.Set as Set
import qualified Language.Haskell.TH.Syntax as TH

import Symantic.Univariant.Trans
import Symantic.Parser.Grammar.Combinators (ErrorItem(..))
import Symantic.Parser.Machine.Input
import Symantic.Parser.Machine.Instructions
import qualified Symantic.Parser.Haskell as H

-- * Type 'Gen'
-- | Generate the 'CodeQ' parsing the input.
newtype Gen inp vs es a = Gen { unGen ::
  GenCtx inp vs es a ->
  CodeQ (Either (ParsingError inp) a)
}

-- ** Type 'ParsingError'
data ParsingError inp
  =  ParsingErrorStandard
  {  parsingErrorOffset :: Offset
  ,  parsingErrorUnexpected :: Maybe (InputToken inp)
  ,  parsingErrorExpecting :: Set (ErrorItem (InputToken inp))
  }
deriving instance Show (InputToken inp) => Show (ParsingError inp)

-- ** Type 'Offset'
type Offset = Int

-- ** Type 'Cont'
type Cont inp v a =
  {-farthestInput-}Cursor inp ->
  {-farthestExpecting-}[ErrorItem (InputToken inp)] ->
  v ->
  Cursor inp ->
  Either (ParsingError inp) a

-- ** Type 'SubRoutine'
type SubRoutine inp v a =
  {-ok-}Cont inp v a ->
  Cursor inp ->
  {-ko-}FailHandler inp a ->
  Either (ParsingError inp) a

-- ** Type 'FailHandler'
type FailHandler inp a =
  {-failureInput-}Cursor inp ->
  {-farthestInput-}Cursor inp ->
  {-farthestExpecting-}[ErrorItem (InputToken inp)] ->
  Either (ParsingError inp) a

{-
-- *** Type 'FarthestError'
data FarthestError inp = FarthestError
  { farthestInput :: Cursor inp
  , farthestExpecting :: [ErrorItem (InputToken inp)]
  }
-}

-- | @('generate' input mach)@ generates @TemplateHaskell@ code
-- parsing given 'input' according to given 'mach'ine.
generate ::
  forall inp ret.
  Ord (InputToken inp) =>
  Show (InputToken inp) =>
  TH.Lift (InputToken inp) =>
  -- InputToken inp ~ Char =>
  Input inp =>
  CodeQ inp ->
  Show (Cursor inp) =>
  Gen inp '[] ('Succ 'Zero) ret ->
  CodeQ (Either (ParsingError inp) ret)
generate input (Gen k) = [||
  -- Pattern bindings containing unlifted types
  -- should use an outermost bang pattern.
  let !(# init, readMore, readNext #) = $$(cursorOf input) in
  let finalRet = \_farInp _farExp v _inp -> Right v in
  let finalFail _failInp !farInp !farExp =
        Left ParsingErrorStandard
        { parsingErrorOffset = offset farInp
        , parsingErrorUnexpected =
            if readMore farInp
            then Just (let (# c, _ #) = readNext farInp in c)
            else Nothing
        , parsingErrorExpecting = Set.fromList farExp
        } in
  $$(k GenCtx
    { valueStack = ValueStackEmpty
    , failStack = FailStackCons [||finalFail||] FailStackEmpty
    , retCode = [||finalRet||]
    , input = [||init||]
    , nextInput = [||readNext||]
    , moreInput = [||readMore||]
    -- , farthestError = [||Nothing||]
    , farthestInput = [||init||]
    , farthestExpecting = [|| [] ||]
    })
  ||]

-- ** Type 'GenCtx'
-- | This is a context only present at compile-time.
data GenCtx inp vs (es::Peano) a =
  ( TH.Lift (InputToken inp)
  , Cursorable (Cursor inp)
  , Show (InputToken inp)
  -- , InputToken inp ~ Char
  ) => GenCtx
  { valueStack :: ValueStack vs
  , failStack :: FailStack inp es a
  , retCode :: CodeQ (Cont inp a a)
  , input :: CodeQ (Cursor inp)
  , moreInput :: CodeQ (Cursor inp -> Bool)
  , nextInput :: CodeQ (Cursor inp -> (# InputToken inp, Cursor inp #))
  , farthestInput :: CodeQ (Cursor inp)
  , farthestExpecting :: CodeQ [ErrorItem (InputToken inp)]
  }

-- ** Type 'ValueStack'
data ValueStack vs where
  ValueStackEmpty :: ValueStack '[]
  ValueStackCons ::
    -- TODO: maybe use H.Haskell instead of CodeQ ?
    -- as in https://github.com/j-mie6/ParsleyHaskell/popFail/3ec0986a5017866919a6404c14fe78678b7afb46
    { valueStackHead :: CodeQ v
    , valueStackTail :: ValueStack vs
    } -> ValueStack (v ': vs)

-- ** Type 'FailStack'
data FailStack inp es a where
  FailStackEmpty :: FailStack inp 'Zero a
  FailStackCons ::
    { failStackHead :: CodeQ (FailHandler inp a)
    , failStackTail :: FailStack inp es a
    } ->
    FailStack inp ('Succ es) a

instance Stackable Gen where
  push x k = Gen $ \ctx -> unGen k ctx
    { valueStack = ValueStackCons (liftCode x) (valueStack ctx) }
  pop k = Gen $ \ctx -> unGen k ctx
    { valueStack = valueStackTail (valueStack ctx) }
  liftI2 f k = Gen $ \ctx -> unGen k ctx
    { valueStack =
      let ValueStackCons y (ValueStackCons x xs) = valueStack ctx in
      ValueStackCons (liftCode2 f x y) xs
    }
  swap k = Gen $ \ctx -> unGen k ctx
    { valueStack =
        let ValueStackCons y (ValueStackCons x xs) = valueStack ctx in
        ValueStackCons x (ValueStackCons y xs)
    }
instance Branchable Gen where
  case_ kx ky = Gen $ \ctx ->
    let ValueStackCons v vs = valueStack ctx in
    [||
      case $$v of
        Left  x -> $$(unGen kx ctx{ valueStack = ValueStackCons [||x||] vs })
        Right y -> $$(unGen ky ctx{ valueStack = ValueStackCons [||y||] vs })
    ||]
  choices fs ks kd = Gen $ \ctx ->
    let ValueStackCons v vs = valueStack ctx in
    go ctx{valueStack = vs} v fs ks
    where
    go ctx x (f:fs') (Gen k:ks') = [||
      if $$(liftCode1 f x) then $$(k ctx)
      else $$(go ctx x fs' ks')
      ||]
    go ctx _ _ _ = unGen kd ctx
instance Failable Gen where
  fail failExp = Gen $ \ctx@GenCtx{} -> [||
    let (# farInp, farExp #) =
          case $$compareOffset $$(farthestInput ctx) $$(input ctx) of
            LT -> (# $$(input ctx), failExp #)
            EQ -> (# $$(farthestInput ctx), ($$(farthestExpecting ctx) <> failExp) #)
            GT -> (# $$(farthestInput ctx), $$(farthestExpecting ctx) #) in
    {-
    trace ("fail: "
      <>" failExp="<>show @[ErrorItem Char] failExp
      <>" farthestExpecting="<>show @[ErrorItem Char] ($$(farthestExpecting ctx))
      <>" farExp="<>show @[ErrorItem Char] farExp) $
    -}
    $$(failStackHead (failStack ctx))
      $$(input ctx) farInp farExp
    ||]
  popFail k = Gen $ \ctx ->
    let FailStackCons _e es = failStack ctx in
    unGen k ctx{failStack = es}
  catchFail ok ko = Gen $ \ctx@GenCtx{} -> [||
    let _ = "catchFail" in $$(unGen ok ctx
      { failStack = FailStackCons [|| \(!failInp) (!farInp) (!farExp) ->
          -- trace ("catchFail: " <> "farExp="<>show farExp) $
          $$(unGen ko ctx
            -- Push the input as it was when entering the catchFail.
            { valueStack = ValueStackCons (input ctx) (valueStack ctx)
            -- Move the input to the failing position.
            , input = [||failInp||]
            -- Set the farthestInput to the farthest computed by 'fail'
            , farthestInput = [||farInp||]
            , farthestExpecting = [||farExp||]
            })
        ||] (failStack ctx)
      })
    ||]
instance Inputable Gen where
  loadInput k = Gen $ \ctx ->
    let ValueStackCons input vs = valueStack ctx in
    unGen k ctx{valueStack = vs, input}
  pushInput k = Gen $ \ctx ->
    unGen k ctx{valueStack = ValueStackCons (input ctx) (valueStack ctx)}
instance Routinable Gen where
  call (LetName n) k = Gen $ \ctx -> [||
    let _ = "call" in
    $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
      $$(suspend k ctx)
      $$(input ctx)
      $! $$(failStackHead (failStack ctx))
    ||]
  jump (LetName n) = Gen $ \ctx -> [||
    let _ = "jump" in
    $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
      $$(retCode ctx)
      $$(input ctx)
      $! $$(failStackHead (failStack ctx))
    ||]
  ret = Gen $ \ctx -> unGen (resume (retCode ctx)) ctx
  subroutine (LetName n) sub k = Gen $ \ctx -> Code $ TH.unsafeTExpCoerce $ do
    body <- TH.unTypeQ $ TH.examineCode $ [|| -- buildRec in Parsley
      -- SubRoutine
      -- Why using $! at call site and not ! here on ko?
      \ !ok !inp ko ->
        $$(unGen sub ctx
          { valueStack = ValueStackEmpty
          , failStack = FailStackCons [||ko||] FailStackEmpty
          , input = [||inp||]
          , retCode = [||ok||]
          -- , farthestInput = [|inp|]
          -- , farthestExpecting = [|| [] ||]
          })
      ||]
    let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
    expr <- TH.unTypeQ (TH.examineCode (unGen k ctx))
    return (TH.LetE [decl] expr)

suspend ::
  {-k-}Gen inp (v ': vs) es a ->
  GenCtx inp vs es a ->
  CodeQ (Cont inp v a)
suspend k ctx = [||
  let _ = "suspend" in
  \farInp farExp v !inp ->
    $$(unGen k ctx
      { valueStack = ValueStackCons [||v||] (valueStack ctx)
      , input = [||inp||]
      , farthestInput = [||farInp||]
      , farthestExpecting = [||farExp||]
      }
    )
  ||]

resume :: CodeQ (Cont inp v a) -> Gen inp (v ': vs) es a
resume k = Gen $ \ctx -> [||
  let _ = "resume" in
  $$k
    $$(farthestInput ctx)
    $$(farthestExpecting ctx)
    $$(valueStackHead (valueStack ctx))
    $$(input ctx)
  ||]

instance Joinable Gen where
  defJoin (LetName n) sub k = Gen $ \ctx -> Code $ TH.unsafeTExpCoerce $ do
    body <- TH.unTypeQ $ TH.examineCode $ [||
      \farInp farExp v !inp ->
        $$(unGen sub ctx
          { valueStack = ValueStackCons [||v||] (valueStack ctx)
          , input = [||inp||]
          , farthestInput = [||farInp||]
          , farthestExpecting = [||farExp||]
          })
      ||]
    let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
    expr <- TH.unTypeQ (TH.examineCode (unGen k ctx))
    return (TH.LetE [decl] expr)
  refJoin (LetName n) =
    resume (Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
instance Readable Gen Char where
  read farExp p k =
    -- TODO: piggy bank
    maybeEmitCheck (Just 1) k
    where
    maybeEmitCheck Nothing ok = sat (liftCode p) ok (fail farExp)
    maybeEmitCheck (Just n) ok = Gen $ \ctx ->
      let FailStackCons e es = failStack ctx in
      [||
      let readFail = $$(e) in -- Factorize failure code
      $$((`unGen` ctx{failStack = FailStackCons [||readFail||] es}) $ emitLengthCheck n
        {-ok-}(sat (liftCode p) ok
          {-ko-}(fail farExp))
        {-ko-}(fail farExp))
      ||]

sat ::
  forall inp vs es a.
  -- Cursorable (Cursor inp) =>
  -- InputToken inp ~ Char =>
  Ord (InputToken inp) =>
  TH.Lift (InputToken inp) =>
  {-predicate-}CodeQ (InputToken inp -> Bool) ->
  {-ok-}Gen inp (InputToken inp ': vs) ('Succ es) a ->
  {-ko-}Gen inp vs ('Succ es) a ->
  Gen inp vs ('Succ es) a
sat p ok ko = Gen $ \ctx -> [||
  let !(# c, cs #) = $$(nextInput ctx) $$(input ctx) in
  if $$p c
  then $$(unGen ok ctx
    { valueStack = ValueStackCons [||c||] (valueStack ctx)
    , input = [||cs||]
    })
  else let _ = "sat.else" in $$(unGen ko ctx)
  ||]

{-
evalSat ::
  -- Cursorable inp =>
  -- HandlerOps inp =>
  InstrPure (Char -> Bool) ->
  Gen inp (Char ': vs) ('Succ es) a ->
  Gen inp vs ('Succ es) a
evalSat p k = do
  bankrupt <- asks isBankrupt
  hasChange <- asks hasCoin
  if | bankrupt -> maybeEmitCheck (Just 1) <$> k
     | hasChange -> maybeEmitCheck Nothing <$> local spendCoin k
     | otherwise -> local breakPiggy (maybeEmitCheck . Just <$> asks coins <*> local spendCoin k)
  where
  maybeEmitCheck Nothing mk ctx = sat (genDefunc p) mk (raise ctx) ctx
  maybeEmitCheck (Just n) mk ctx =
    [|| let bad = $$(raise ctx) in $$(emitLengthCheck n (sat (genDefunc p) mk [||bad||]) [||bad||] ctx)||]
-}

emitLengthCheck ::
  TH.Lift (InputToken inp) =>
  Int -> Gen inp vs es a -> Gen inp vs es a -> Gen inp vs es a
emitLengthCheck 0 ok _ko = ok
emitLengthCheck 1 ok ko = Gen $ \ctx -> [||
  if $$(moreInput ctx) $$(input ctx)
  then $$(unGen ok ctx)
  else let _ = "sat.length-check.else" in $$(unGen ko ctx)
  ||]
{-
emitLengthCheck n ok ko ctx = Gen $ \ctx -> [||
  if $$moreInput ($$shiftRight $$(input ctx) (n - 1))
  then $$(unGen ok ctx)
  else $$(unGen ko ctx {farthestExpecting = [||farExp||]})
  ||]
-}


liftCode :: InstrPure a -> CodeQ a
liftCode = trans
{-# INLINE liftCode #-}

liftCode1 :: InstrPure (a -> b) -> CodeQ a -> CodeQ b
liftCode1 p a = case p of
  InstrPureSameOffset -> [|| $$sameOffset $$a ||]
  InstrPureHaskell h -> go a h
  where
  go :: CodeQ a -> H.Haskell (a -> b) -> CodeQ b
  go qa = \case
    (H.:$) -> [|| \x -> $$qa x ||]
    (H.:.) -> [|| \g x -> $$qa (g x) ||]
    H.Flip -> [|| \x y -> $$qa y x ||]
    (H.:.) H.:@ f H.:@ g -> [|| $$(go (go qa g) f) ||]
    H.Const -> [|| \_ -> $$qa ||]
    H.Flip H.:@ H.Const -> H.id
    h@(H.Flip H.:@ _f) -> [|| \x -> $$(liftCode2 (InstrPureHaskell h) qa [||x||]) ||]
    H.Eq x -> [|| $$(trans x) Eq.== $$qa ||]
    H.Id -> qa
    h -> [|| $$(trans h) $$qa ||]

liftCode2 :: InstrPure (a -> b -> c) -> CodeQ a -> CodeQ b -> CodeQ c
liftCode2 p a b = case p of
  InstrPureSameOffset -> [|| $$sameOffset $$a $$b ||]
  InstrPureHaskell h -> go a b h
  where
  go :: CodeQ a -> CodeQ b -> H.Haskell (a -> b -> c) -> CodeQ c
  go qa qb = \case
    (H.:$) -> [|| $$qa $$qb ||]
    (H.:.) -> [|| \x -> $$qa ($$qb x) ||]
    H.Flip -> [|| \x -> $$qa x $$qb ||]
    H.Flip H.:@ H.Const -> [|| $$qb ||]
    H.Flip H.:@ f -> go qb qa f
    H.Const -> [|| $$qa ||]
    H.Cons -> [|| $$qa : $$qb ||]
    h -> [|| $$(trans h) $$qa $$qb ||]