{-# 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.List (minimum)
import Data.Map (Map)
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.Map.Strict as Map
import qualified Data.Set as Set
import qualified Language.Haskell.TH.Syntax as TH
-- import qualified Control.Monad.Trans.Writer as Writer

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

-- * Type 'Gen'
-- | Generate the 'CodeQ' parsing the input.
data Gen inp vs es a = Gen
  { minHorizon :: Map TH.Name Horizon -> Horizon
  , unGen ::
      GenCtx inp vs es a ->
      CodeQ (Either (ParsingError inp) a)
  }

-- ** Type 'ParsingError'
data ParsingError inp
  =  ParsingErrorStandard
  {  parsingErrorOffset :: Offset
    -- | Note that if an 'ErrorItemHorizon' greater than 1
    -- is amongst the 'parsingErrorExpecting'
    -- then this is only the 'InputToken'
    -- at the begining of the expected 'Horizon'.
  ,  parsingErrorUnexpected :: Maybe (InputToken inp)
  ,  parsingErrorExpecting :: Set (ErrorItem (InputToken inp))
  }
deriving instance Show (InputToken inp) => Show (ParsingError inp)

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

-- ** Type 'Horizon'
-- | Synthetized minimal input length
-- required for a successful parsing.
-- Used with 'horizon' to factorize input length checks,
-- instead of checking the input length
-- one 'InputToken' by one 'InputToken' at each 'read'.
type Horizon = Offset

-- ** 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 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
  $$(unGen k GenCtx
    { valueStack = ValueStackEmpty
    , failStack = FailStackCons [||finalFail||] FailStackEmpty
    , retCode = [||finalRet||]
    , input = [||init||]
    , nextInput = [||readNext||]
    , moreInput = [||readMore||]
    -- , farthestError = [||Nothing||]
    , farthestInput = [||init||]
    , farthestExpecting = [|| [] ||]
    , horizon = 0
    , horizonByName = Map.empty
    })
  ||]

-- ** 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)]
    -- | Remaining horizon
  , horizon :: Offset
    -- | Horizon for each 'call' or 'jump'.
  , horizonByName :: Map TH.Name Offset
  }

-- ** Type 'ValueStack'
data ValueStack vs where
  ValueStackEmpty :: ValueStack '[]
  ValueStackCons ::
    -- TODO: maybe use H.CombPure 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 = k
    { unGen = \ctx -> unGen k ctx
      { valueStack = ValueStackCons (liftCode x) (valueStack ctx) }
    }
  pop k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack = valueStackTail (valueStack ctx) }
    }
  liftI2 f k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack =
        let ValueStackCons y (ValueStackCons x xs) = valueStack ctx in
        ValueStackCons (liftCode2 f x y) xs
      }
    }
  swap k = k
    { unGen = \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
    { minHorizon = \ls ->
      minHorizon kx ls `min` minHorizon ky ls
    , unGen = \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
    { minHorizon = \ls -> minimum $
        minHorizon kd ls :
        (($ ls) . minHorizon <$> ks)
    , unGen = \ctx ->
      let ValueStackCons v vs = valueStack ctx in
      go ctx{valueStack = vs} v fs ks
    }
    where
    go ctx x (f:fs') (k:ks') = [||
      if $$(liftCode1 f x) then $$(unGen k ctx)
      else $$(go ctx x fs' ks')
      ||]
    go ctx _ _ _ = unGen kd ctx
instance Failable Gen where
  fail failExp = Gen
    { minHorizon = \_hs -> 0
    , unGen = \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
      $$(failStackHead (failStack ctx))
        $$(input ctx) farInp farExp
      ||]
    }
  popFail k = k
    { unGen = \ctx ->
      let FailStackCons _e es = failStack ctx in
      unGen k ctx{failStack = es}
    }
  catchFail ok ko = Gen
    { minHorizon = \ls -> minHorizon ok ls `min` minHorizon ko ls
    , unGen = \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 = k
    { unGen = \ctx ->
      let ValueStackCons input vs = valueStack ctx in
      unGen k ctx
        { valueStack = vs
        , input
        , horizon = 0
        }
    }
  pushInput k = k
    { unGen = \ctx ->
      unGen k ctx{valueStack = ValueStackCons (input ctx) (valueStack ctx)}
    }
instance Routinable Gen where
  call (LetName n) k = k
    { minHorizon = \hs -> hs Map.! n
    , unGen = \ctx -> [||
      let _ = "call" in
      $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
        {-ok-}$$(generateSuspend k ctx)
        $$(input ctx)
        $! $$(failStackHead (failStack ctx))
      ||]
    }
  jump (LetName n) = Gen
    { minHorizon = \hs -> hs Map.! n
    , unGen = \ctx -> [||
      let _ = "jump" in
      $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
        {-ok-}$$(retCode ctx)
        $$(input ctx)
        $! $$(failStackHead (failStack ctx))
      ||]
    }
  ret = Gen
    { minHorizon = \_hs -> 0
    , unGen = \ctx -> unGen (generateResume (retCode ctx)) ctx
    }
  subroutine (LetName n) sub k = Gen
    { minHorizon = \hs ->
        minHorizon k $
          Map.insert n (minHorizon sub (Map.insert n 0 hs)) hs
    , unGen = \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 = [|| [] ||]
            , horizon = 0
            , horizonByName = Map.insert n 0 (horizonByName ctx)
            })
        ||]
      let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
      expr <- TH.unTypeQ (TH.examineCode (unGen k ctx
        { horizonByName =
            Map.insert n
              (minHorizon sub
                (Map.insert n 0 (horizonByName ctx)))
              (horizonByName ctx)
        }))
      return (TH.LetE [decl] expr)
    }

-- | Generate a continuation to be called with 'generateResume',
-- used when 'call' 'ret'urns.
generateSuspend ::
  {-k-}Gen inp (v ': vs) es a ->
  GenCtx inp vs es a ->
  CodeQ (Cont inp v a)
generateSuspend k ctx = [||
  let _ = "suspend" in
  \farInp farExp v !inp ->
    $$(unGen k ctx
      { valueStack = ValueStackCons [||v||] (valueStack ctx)
      , input = [||inp||]
      , farthestInput = [||farInp||]
      , farthestExpecting = [||farExp||]
      , horizon = 0
      }
    )
  ||]

-- | Generate a call to the 'generateSuspend' continuation,
-- used when 'call' 'ret'urns.
generateResume ::
  CodeQ (Cont inp v a) ->
  Gen inp (v ': vs) es a
generateResume k = Gen
  { minHorizon = \_hs -> 0
  , unGen = \ctx -> [||
    let _ = "resume" in
    $$k
      $$(farthestInput ctx)
      $$(farthestExpecting ctx)
      $$(valueStackHead (valueStack ctx))
      $$(input ctx)
    ||]
  }

instance Joinable Gen where
  defJoin (LetName n) sub k = k
    { minHorizon = \hs ->
        minHorizon k $
          Map.insert n (minHorizon sub (Map.insert n 0 hs)) hs
    , unGen = \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||]
            , horizon = 0
            , horizonByName = Map.insert n 0 (horizonByName ctx)
            })
        ||]
      let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
      expr <- TH.unTypeQ (TH.examineCode (unGen k ctx
        { horizonByName =
            Map.insert n
              (minHorizon sub
                (Map.insert n 0 (horizonByName ctx)))
              (horizonByName ctx)
        }))
      return (TH.LetE [decl] expr)
    }
  refJoin (LetName n) =
    generateResume (Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
instance Readable Gen Char where
  read farExp p = checkHorizon . checkToken farExp (liftCode p)

checkHorizon ::
  TH.Lift (InputToken inp) =>
  {-ok-}Gen inp vs ('Succ es) a ->
  Gen inp vs ('Succ es) a
checkHorizon ok = ok
  { minHorizon = \hs -> 1 + minHorizon ok hs
  , unGen = \ctx0@GenCtx{failStack = FailStackCons e es} -> [||
      -- Factorize failure code
      let readFail = $$(e) in
      $$(
        let ctx = ctx0{ failStack = FailStackCons [||readFail||] es } in
        if horizon ctx >= 1
        then unGen ok ctx0{horizon = horizon ctx - 1}
        else let minHoz = minHorizon ok (horizonByName ctx) in
          [||
          if $$(moreInput ctx)
               $$(if minHoz > 0
                 then [||$$shiftRight minHoz $$(input ctx)||]
                 else input ctx)
          then $$(unGen ok ctx{horizon = minHoz})
          else let _ = "checkHorizon.else" in
            $$(unGen (fail [ErrorItemHorizon (minHoz + 1)]) ctx)
          ||]
      )
    ||]
  }

checkToken ::
  forall inp vs es a.
  Ord (InputToken inp) =>
  TH.Lift (InputToken inp) =>
  [ErrorItem (InputToken inp)] ->
  {-predicate-}CodeQ (InputToken inp -> Bool) ->
  {-ok-}Gen inp (InputToken inp ': vs) ('Succ es) a ->
  Gen inp vs ('Succ es) a
checkToken farExp p ok = ok
  { unGen = \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 _ = "checkToken.else" in $$(unGen (fail farExp) ctx)
    ||]
  }

liftCode :: InstrPure a -> CodeQ a
liftCode x = trans x
{-
liftCode p = case p of
  InstrPureSameOffset -> [|| $$sameOffset ||]
  InstrPure h -> go h
  where
  go :: H.CombPure a -> CodeQ a
  go = \case
    ((H.:.) H.:@ f) H.:@ (H.Const H.:@ x) -> [|| $$(go f) $$(go x) ||]
    a -> trans a
-}
-- {-# INLINE liftCode #-}

liftCode1 :: InstrPure (a -> b) -> CodeQ a -> CodeQ b
liftCode1 p a = case p of
  InstrPureSameOffset f -> [|| $$f $$a ||]
  InstrPure h -> go a h
  where
  go :: CodeQ a -> H.CombPure (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.:@ (H.Const H.@ x) -> [|| $$(go (go qa g) f) ||]
    (H.:.) H.:@ f H.:@ g -> [|| $$(go (go qa g) f) ||]
    H.Cons -> [|| ($$qa :) ||]
    H.Const -> [|| \_ -> $$qa ||]
    H.Flip H.:@ H.Const -> H.id
    h@(H.Flip H.:@ _f) -> [|| \x -> $$(liftCode2 (InstrPure h) qa [||x||]) ||]
    H.Id H.:@ x -> go qa x
    H.Eq x -> [|| $$(trans x) Eq.== $$qa ||]
    H.Id -> qa
    H.CombPure (H.ValueCode _a2b qa2b) -> [|| $$qa2b $$qa ||]
    -- h -> [|| $$(liftCode h) $$qa ||]

liftCode2 :: InstrPure (a -> b -> c) -> CodeQ a -> CodeQ b -> CodeQ c
liftCode2 p a b = case p of
  InstrPureSameOffset f -> [|| $$f $$a $$b ||]
  InstrPure h -> go a b h
  where
  go :: CodeQ a -> CodeQ b -> H.CombPure (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.Id H.:@ x -> go qa qb x
    H.Id -> [|| $$qa $$qb ||]
    H.Cons -> [|| $$qa : $$qb ||]
    H.Const -> [|| $$qa ||]
    H.CombPure (H.ValueCode _a2b2c qa2b2c) -> [|| $$qa2b2c $$qa $$qb ||]
    --h -> [|| $$(trans h) $$qa $$qb ||]