{-# 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.List.NonEmpty (NonEmpty(..))
import Data.Map (Map)
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..), Ordering(..))
import Data.Semigroup (Semigroup(..))
import Data.Set (Set)
import Data.String (String)
import Language.Haskell.TH (CodeQ, Code(..))
import Prelude ((+), (-))
import Text.Show (Show(..))
import GHC.TypeLits (symbolVal)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Internal as Map_
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Language.Haskell.TH as TH
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

import Debug.Trace (trace)

genCode :: TermInstr a -> CodeQ a
genCode = trans

-- * Type 'Gen'
-- | Generate the 'CodeQ' parsing the input.
data Gen inp vs a = Gen
  { minHorizon :: Map TH.Name Horizon -> Horizon
    -- ^ Synthetized (bottom-up) minimal input length
    -- required by the parser to not fail.
    -- This requires a 'minHorizonByName'
    -- containing the minimal 'Horizon's of all the 'TH.Name's
    -- this parser 'call's, 'jump's or 'refJoin's to.
  , exceptions :: Map TH.Name (Map ErrorLabel ()) -> Map ErrorLabel ()
  , unGen ::
      GenCtx inp vs 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 'ErrorLabel'
type ErrorLabel = String

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

-- ** Type 'Horizon'
-- | Synthetized minimal input length
-- required for a successful parsing.
-- Used with 'checkedHorizon' to factorize input length checks,
-- instead of checking the input length
-- one 'InputToken' at a time 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 'FarthestError'
data FarthestError inp = FarthestError
  { farthestInput :: Cursor inp
  , farthestExpecting :: [ErrorItem (InputToken inp)]
  }
-}

-- | @('generateCode' input mach)@ generates @TemplateHaskell@ code
-- parsing the given 'input' according to the given 'Machine'.
generateCode ::
  Ord (InputToken inp) =>
  Show (InputToken inp) =>
  TH.Lift (InputToken inp) =>
  -- InputToken inp ~ Char =>
  Input inp =>
  Show (Cursor inp) =>
  Gen inp '[] a ->
  CodeQ (inp -> Either (ParsingError inp) a)
generateCode k = [|| \(input :: inp) ->
  -- 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 finalRaise :: forall b. (Catcher inp b)
        = \_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
    , catchStackByLabel = Map.empty
    , defaultCatch = [||finalRaise||]
    , retCode = [||finalRet||]
    , input = [||init||]
    , nextInput = [||readNext||]
    , moreInput = [||readMore||]
    -- , farthestError = [||Nothing||]
    , farthestInput = [||init||]
    , farthestExpecting = [|| [] ||]
    , checkedHorizon = 0
    , minHorizonByName = Map.empty
    , exceptionsByName = Map.empty
    })
  ||]

-- ** Type 'GenCtx'
-- | This is an inherited (top-down) context
-- only present at compile-time, to build TemplateHaskell splices.
data GenCtx inp vs a =
  ( TH.Lift (InputToken inp)
  , Cursorable (Cursor inp)
  , Show (InputToken inp)
  ) => GenCtx
  { valueStack :: ValueStack vs
  , catchStackByLabel :: Map ErrorLabel (NonEmpty (CodeQ (Catcher inp a)))
    -- | Default 'Catcher' defined at the begining of the generated 'CodeQ',
    -- hence a constant within the 'Gen'eration.
  , defaultCatch :: forall b. CodeQ (Catcher inp b)
  , 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 already checked.
    -- Updated by 'checkHorizon'
    -- and reset elsewhere when needed.
  , checkedHorizon :: Horizon
    -- | Minimal horizon for each 'defLet' or 'defJoin'.
    -- This can be done as an inherited attribute because
    -- 'OverserveSharing' introduces 'def' as an ancestor node
    -- of all the 'ref's pointing to it.
    -- Same for 'defJoin' and its 'refJoin's.
  , minHorizonByName :: Map TH.Name Horizon
  , exceptionsByName :: Map TH.Name (Map ErrorLabel ())
  }

-- ** Type 'ValueStack'
data ValueStack vs where
  ValueStackEmpty :: ValueStack '[]
  ValueStackCons ::
    { valueStackHead :: TermInstr v
    , valueStackTail :: ValueStack vs
    } -> ValueStack (v ': vs)

instance InstrValuable Gen where
  pushValue x k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack = ValueStackCons x (valueStack ctx) }
    }
  popValue k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack = valueStackTail (valueStack ctx) }
    }
  lift2Value f k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack =
        let ValueStackCons y (ValueStackCons x xs) = valueStack ctx in
        ValueStackCons (f H.:@ x H.:@ y) xs
      }
    }
  swapValue k = k
    { unGen = \ctx -> unGen k ctx
      { valueStack =
          let ValueStackCons y (ValueStackCons x xs) = valueStack ctx in
          ValueStackCons x (ValueStackCons y xs)
      }
    }
instance InstrBranchable Gen where
  caseBranch kx ky = Gen
    { minHorizon = \hs -> minHorizon kx hs `min` minHorizon ky hs
    , exceptions = \hs -> exceptions kx hs <> exceptions ky hs
    , unGen = \ctx ->
      let ValueStackCons v vs = valueStack ctx in
      [||
        case $$(genCode v) of
          Left  x -> $$(unGen kx ctx{ valueStack = ValueStackCons (H.Term [||x||]) vs })
          Right y -> $$(unGen ky ctx{ valueStack = ValueStackCons (H.Term [||y||]) vs })
      ||]
    }
  choicesBranch fs ks kd = Gen
    { minHorizon = \hs -> minimum $
        minHorizon kd hs :
        (($ hs) . minHorizon <$> ks)
    , exceptions = \hs -> mconcat $
        exceptions kd hs :
        (($ hs) . exceptions <$> 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 $$(genCode (f H.:@ x))
      then $$(unGen k ctx)
      else $$(go ctx x fs' ks')
      ||]
    go ctx _ _ _ = unGen kd ctx
instance InstrExceptionable Gen where
  raiseException lbl failExp = Gen
    { minHorizon = \_hs -> 0
    , exceptions = \_hs -> Map.singleton (symbolVal lbl) ()
    , 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
        $$(NE.head (Map.findWithDefault (NE.singleton (defaultCatch ctx)) (symbolVal lbl) (catchStackByLabel ctx)))
        $$(input ctx) farInp farExp
      ||]
    }
  popException lbl k = k
    { unGen = \ctx ->
      unGen k ctx{catchStackByLabel = Map.update (\case
          _r0:|(r1:rs) -> Just (r1:|rs)
          _ -> Nothing
        ) (symbolVal lbl) (catchStackByLabel ctx)
      }
    }
  catchException lbl ok ko = Gen
    { minHorizon = \hs -> minHorizon ok hs `min` minHorizon ko hs
    , exceptions = \hs -> exceptions ok hs <> exceptions ko hs
    , unGen = \ctx@GenCtx{} -> [||
        let _ = $$(TH.liftTyped ("catchException lbl="<>symbolVal lbl)) in
        $$(unGen ok ctx
        { catchStackByLabel = Map.insertWith (<>) (symbolVal lbl)
            (NE.singleton ([|| \ !failInp !farInp !farExp ->
              $$(unGen ko ctx
                -- PushValue the input as it was when entering the catchFail.
                { valueStack = ValueStackCons (H.Term (input ctx)) (valueStack ctx)
                -- Note that 'catchStackByLabel' is reset.
                -- Move the input to the failing position.
                , input = [||failInp||]
                -- Set the farthestInput to the farthest computed by 'fail'
                , farthestInput = [||farInp||]
                , farthestExpecting = [||farExp||]
                })
            ||])) (catchStackByLabel ctx)
        }
      ) ||]
    }
-- ** Type 'Catcher'
type Catcher inp a =
  {-failureInput-}Cursor inp ->
  {-farthestInput-}Cursor inp ->
  {-farthestExpecting-}[ErrorItem (InputToken inp)] ->
  Either (ParsingError inp) a
instance InstrInputable Gen where
  loadInput k = k
    { unGen = \ctx ->
      let ValueStackCons input vs = valueStack ctx in
      unGen k ctx
        { valueStack = vs
        , input = genCode input
        , checkedHorizon = 0
        }
    }
  pushInput k = k
    { unGen = \ctx ->
      unGen k ctx{valueStack = ValueStackCons (H.Term (input ctx)) (valueStack ctx)}
    }
instance InstrLetable Gen where
  defLet (LetName n) sub k = k
    { unGen = \ctx@GenCtx{} -> Code $ TH.unsafeTExpCoerce $ do
      -- 'sub' is recursively 'call'able within 'sub',
      -- but its maximal 'minHorizon' is not known yet.
      let minHorizonByNameButSub = Map.insert n 0 (minHorizonByName ctx)
      let raiseLabelsByNameButSub = Map.insert n Map.empty (exceptionsByName ctx)
      body <- TH.unTypeQ $ TH.examineCode $ [|| -- buildRec in Parsley
        -- Called by 'call' or 'jump'.
        \ !ok{-from generateSuspend or retCode-}
          !inp
          !koByLabel{- 'catchStackByLabel' from the 'call'-site -} ->
          $$(unGen sub ctx
            { valueStack = ValueStackEmpty
            -- Build a 'catchStackByLabel' from the one available at the 'call'-site.
            -- Note that all the 'exceptions' of the 'sub'routine may not be available,
            -- hence 'Map.findWithDefault' is used instead of 'Map.!'.
            , catchStackByLabel = Map.mapWithKey
                (\lbl () -> NE.singleton [||Map.findWithDefault $$(defaultCatch ctx) lbl koByLabel||])
                (exceptions sub raiseLabelsByNameButSub)
            , input = [||inp||]
            , retCode = [||ok||]

            -- These are passed by the caller via 'ok' or 'ko'
            -- , farthestInput = 
            -- , farthestExpecting = 

            -- Some callers can call this 'defLet'
            -- with zero 'checkedHorizon', hence use this minimum.
            -- TODO: maybe it could be improved a bit
            -- by taking the minimum of the checked horizons
            -- before all the 'call's and 'jump's to this 'defLet'.
            , checkedHorizon = 0
            , minHorizonByName = minHorizonByNameButSub
            , exceptionsByName = raiseLabelsByNameButSub
            })
        ||]
      let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
      expr <- TH.unTypeQ (TH.examineCode (unGen k ctx
        { minHorizonByName =
            -- 'sub' is 'call'able within 'k'.
            Map.insert n
              (minHorizon sub minHorizonByNameButSub)
              (minHorizonByName ctx)
        , exceptionsByName =
            Map.insert n
              (exceptions sub raiseLabelsByNameButSub)
              (exceptionsByName ctx)
        }))
      return (TH.LetE [decl] expr)
    }
  jump (LetName n) = Gen
    { minHorizon = (Map.! n)
    , exceptions = (Map.! n)
    , unGen = \ctx -> [||
      let _ = "jump" in
      $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
        {-ok-}$$(retCode ctx)
        $$(input ctx)
        $$(liftTypedRaiseByLabel $
          catchStackByLabel ctx
          -- Pass only the labels raised by the 'defLet'.
          `Map.intersection`
          (exceptionsByName ctx Map.! n)
          )
      ||]
    }
  call (LetName n) k = k
    { minHorizon = (Map.! n)
    , exceptions = (Map.! n)
    , unGen = \ctx -> let ks = (Map.keys (catchStackByLabel ctx)) in [||
      let _ = $$(TH.liftTyped $ "call exceptionsByName("<>show n<>")="<>show (Map.keys (Map.findWithDefault Map.empty n (exceptionsByName ctx))) <> " catchStackByLabel(ctx)="<> show ks) in
      $$(Code (TH.unsafeTExpCoerce (return (TH.VarE n))))
        {-ok-}$$(generateSuspend k ctx)
        $$(input ctx)
        $$(liftTypedRaiseByLabel $
          catchStackByLabel ctx
          -- Pass only the labels raised by the 'defLet'.
          `Map.intersection`
          (exceptionsByName ctx Map.! n)
        )
      ||]
    }
  ret = Gen
    { minHorizon = \_hs -> 0
    , exceptions = \_hs -> Map.empty
    , unGen = \ctx -> unGen (generateResume (retCode ctx)) ctx
    }

-- | Like 'TH.liftTyped' but adjusted to work on 'catchStackByLabel'
-- which already contains 'CodeQ' terms.
-- Moreover, only the 'Catcher' at the top of the stack
-- is needed and thus generated in the resulting 'CodeQ'.
--
-- TODO: Use an 'Array' instead of a 'Map'?
liftTypedRaiseByLabel :: TH.Lift k => Map k (NonEmpty (CodeQ a)) -> CodeQ (Map k a)
liftTypedRaiseByLabel Map_.Tip = [|| Map_.Tip ||]
liftTypedRaiseByLabel (Map_.Bin s k (h:|_hs) l r) =
  [|| Map_.Bin s k $$h $$(liftTypedRaiseByLabel l) $$(liftTypedRaiseByLabel r) ||]

-- | Generate a 'retCode' 'Cont'inuation to be called with 'generateResume'.
-- Used when 'call' 'ret'urns.
-- The return 'v'alue is 'pushValue'ed on the 'valueStack'.
generateSuspend ::
  {-k-}Gen inp (v ': vs) a ->
  GenCtx inp vs a ->
  CodeQ (Cont inp v a)
generateSuspend k ctx = [||
  let _ = $$(TH.liftTyped $ "suspend raiseException=" <> show (exceptionsByName ctx)) in
  \farInp farExp v !inp ->
    $$(unGen k ctx
      { valueStack = ValueStackCons (H.Term [||v||]) (valueStack ctx)
      , input = [||inp||]
      , farthestInput = [||farInp||]
      , farthestExpecting = [||farExp||]
      , checkedHorizon = 0
      }
    )
  ||]

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

instance InstrJoinable Gen where
  defJoin (LetName n) joined k = k
    { unGen = \ctx -> Code $ TH.unsafeTExpCoerce $ do
      body <- TH.unTypeQ $ TH.examineCode $ [||
        -- Called by 'generateResume'.
        \farInp farExp v !inp ->
          $$(unGen joined ctx
            { valueStack = ValueStackCons (H.Term [||v||]) (valueStack ctx)
            , input = [||inp||]
            , farthestInput = [||farInp||]
            , farthestExpecting = [||farExp||]
            , checkedHorizon = 0
            {- FIXME:
            , catchStackByLabel = Map.mapWithKey
                (\lbl () -> NE.singleton [||koByLabel Map.! lbl||])
                (exceptions joined raiseLabelsByNameButSub)
            -}
            })
        ||]
      let decl = TH.FunD n [TH.Clause [] (TH.NormalB body) []]
      expr <- TH.unTypeQ (TH.examineCode (unGen k ctx
        { minHorizonByName =
            -- 'joined' is 'refJoin'able within 'k'.
            Map.insert n
              -- By definition (in 'joinNext')
              -- 'joined' is not recursively 'refJoin'able within 'joined',
              -- hence no need to prevent against recursivity
              -- as has to be done in 'defLet'.
              (minHorizon joined (minHorizonByName ctx))
              (minHorizonByName ctx)
        , exceptionsByName =
            Map.insert n
              (exceptions joined (exceptionsByName ctx))
              (exceptionsByName ctx)
        }))
      return (TH.LetE [decl] expr)
    }
  refJoin (LetName n) = (generateResume (Code (TH.unsafeTExpCoerce (return (TH.VarE n)))))
    { minHorizon = (Map.! n)
    , exceptions = (Map.! n)
    }
instance InstrReadable Char Gen where
  read farExp p = checkHorizon . checkToken farExp p

checkHorizon ::
  TH.Lift (InputToken inp) =>
  {-ok-}Gen inp vs a ->
  Gen inp vs a
checkHorizon ok = ok
  { minHorizon = \hs -> 1 + minHorizon ok hs
  , exceptions = \hs -> Map.insert "fail" () $ exceptions ok hs
  , unGen = \ctx0@GenCtx{} ->
    let raiseByLbl =
          NE.head (Map.findWithDefault (NE.singleton (defaultCatch ctx0)) "fail" (catchStackByLabel ctx0)) in
    [||
      -- Factorize failure code
      let readFail = $$(raiseByLbl) in
      $$(
        let ctx = ctx0{catchStackByLabel = Map.adjust (\(_r:|rs) -> [||readFail||] :| rs) "fail" (catchStackByLabel ctx0)} in
        if checkedHorizon ctx >= 1
        then unGen ok ctx0{checkedHorizon = checkedHorizon ctx - 1}
        else let minHoriz = minHorizon ok (minHorizonByName ctx) in
          [||
          if $$(moreInput ctx)
               $$(if minHoriz > 0
                 then [||$$shiftRight minHoriz $$(input ctx)||]
                 else input ctx)
          then $$(unGen ok ctx{checkedHorizon = minHoriz})
          else let _ = "checkHorizon.else" in
            -- TODO: return a resuming continuation (eg. Partial)
            $$(unGen (fail [ErrorItemHorizon (minHoriz + 1)]) ctx)
          ||]
      )
    ||]
  }

checkToken ::
  Ord (InputToken inp) =>
  TH.Lift (InputToken inp) =>
  [ErrorItem (InputToken inp)] ->
  {-predicate-}TermInstr (InputToken inp -> Bool) ->
  {-ok-}Gen inp (InputToken inp ': vs) a ->
  Gen inp vs a
checkToken farExp p ok = ok
  { exceptions = \hs -> Map.insert "fail" () $ exceptions ok hs
  , unGen = \ctx -> [||
    let !(# c, cs #) = $$(nextInput ctx) $$(input ctx) in
    if $$(genCode p) c
    then $$(unGen ok ctx
      { valueStack = ValueStackCons (H.Term [||c||]) (valueStack ctx)
      , input = [||cs||]
      })
    else let _ = "checkToken.else" in $$(unGen (fail farExp) ctx)
    ||]
  }