{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}

module Generics.Annotations.ParsecYield (WithBounds, AnyWithBounds, Bounds(..), getLeftBounds, getRightBounds, pYield, unit, chainl) where

import Control.Monad.Trans
import Control.Applicative

import qualified Text.Parsec as P
import qualified Text.Parsec.Pos as PP
import Generics.MultiRec
import Generics.Annotations.Annotations

import qualified Generics.Annotations.Yield as Y
import Generics.Annotations.Token

-- | The fixpoint of a data family @s@ recursively annotated with 'Bounds'.
type WithBounds  s = Ann  Bounds s

-- | The fixpoint of a data family @s@ recursively annotated with 'Bounds', with existentially quantified top-level index.
type AnyWithBounds s = AnyAnn Bounds s

instance (Y.MonadYield m) => Y.MonadYield (P.ParsecT s u m) where
  type Y.YieldFam (P.ParsecT s u m) = Y.YieldFam m
  type Y.AnnType  (P.ParsecT s u m) = Y.AnnType m
  yield p ann x = lift (Y.yield p ann x)

-- | Bounds of a subexpression. The outer bounds exclude whitespace while the inner bounds don't.
data Bounds = Bounds { outerBounds :: Range, innerBounds :: Range }
  deriving (Eq, Show)

-- | @(left, right)@
type Range = (Int, Int)

-- | @(leftOuter, leftInner)@
type LeftBounds  = (Int, Int)

-- | @(rightInner, rightOuter)@
type RightBounds = (Int, Int)

-- | Creates a 'Bounds' object from the 'LeftBounds' and 'RightBounds'.
mkBounds :: LeftBounds -> RightBounds -> Bounds
mkBounds (leftOuter, leftInner) (rightInner, rightOuter) =
  Bounds (leftOuter, rightOuter) (leftInner, rightInner)

-- | Yields the left outer and left inner bound of the to be parser subexpression.
getLeftBounds :: (Show t, P.Stream s m t, Symbol t) => P.ParsecT s u m (Int, Int)
getLeftBounds = do
  leftOuter <- getOffset
  Just la <- lookAhead1
  let leftInner = leftOuter + spaceBeforeSize la
  return (leftOuter, leftInner)

-- | Yields the right inner and right outer bound of the just parsed subexpression.
getRightBounds :: (Show t, P.Stream s m t, Symbol t) => P.ParsecT s u m (Int, Int)
getRightBounds = do
  rightInner <- getOffset
  mla <- lookAhead1
  let rightOuter = case mla of
        Nothing -> rightInner
        Just la -> rightInner + spaceBeforeSize la
  return (rightInner, rightOuter)

lookAhead1 :: (Show t, P.Stream s m t) => P.ParsecT s u m (Maybe t)
lookAhead1 = P.lookAhead (P.optionMaybe P.anyToken)

getOffset :: Monad m => P.ParsecT s u m Int
getOffset = (pred . PP.sourceColumn) `fmap` P.getPosition

-- | This specific version of yield annotates the result with @Bounds@.
--   It retrieves the right bounds itself; the left bounds are passed as argument.
pYield :: forall s u m t a. (P.Stream s m t, Symbol t, Show t, Y.MonadYield m, Y.AnnType m ~ Bounds) =>
  Y.YieldFam m a -> LeftBounds -> a -> P.ParsecT s u m a
pYield p lb x = do
  rb <- getRightBounds
  Y.yield p (mkBounds lb rb) x

-- | A combinator for encoding left-recursion.
chainl :: (Y.AnnType m ~ Bounds, Symbol t, Show t, P.Stream s m t, El (Y.YieldFam m) a, Y.MonadYield m) =>
          P.ParsecT s u m a -> P.ParsecT s u m (a -> a -> a) -> P.ParsecT s u m a
chainl px pf = do
    m <- getLeftBounds
    let rest x = option x $ do
          f <- pf
          y <- px
          let z = f x y
          pYield proof m z
          rest z
    x <- px
    rest x

-- | Mark before and yield after running the provided parser.
unit :: (Y.AnnType m ~ Bounds, El (Y.YieldFam m) b, Show t, Symbol t, P.Stream s m t, Y.MonadYield m) =>
          P.ParsecT s u m b -> P.ParsecT s u m b
unit px = do
  m <- getLeftBounds
  x <- px
  pYield proof m x

option :: Alternative f => a -> f a -> f a
option x p = p <|> pure x