{-# LANGUAGE TupleSections #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DerivingVia #-}
module Presyntax.Parser where

import Control.Applicative
import Control.Monad.State

import qualified Data.Text as T
import Data.Text (Text)

import Presyntax.Lexer
import Presyntax

import Syntax


data ParseError
  = UnexpectedEof Int Int
  | Unexpected Token
  | Empty
  | AltError ParseError ParseError
  deriving (Show)

data ParseState
  = ParseState { ptTs   :: [Token]
               , ptLine :: !Int
               , ptCol  :: !Int
               }

newtype Parser a =
  Parser { runParser :: ParseState -> Either ParseError (a, ParseState) }
  deriving
    ( Functor
    , Applicative
    , Monad
    , MonadState ParseState
    )
  via (StateT ParseState (Either ParseError))

parseString :: Parser a -> String -> Either (Either LexError ParseError) a
parseString (Parser k) s =
  case lexString s of
    Left e -> Left (Left e)
    Right xs ->
      case k (ParseState xs 0 0) of
        Left e -> Left (pure e)
        Right (x, _) -> Right x

selectToken :: (Token -> Maybe a) -> Parser a
selectToken k = Parser \case
  ParseState [] l c -> Left (UnexpectedEof l c)
  ParseState (x:xs) _ _ ->
    case k x of
      Just p -> pure (p, ParseState xs (tokLine x) (tokCol x))
      Nothing -> Left (Unexpected x)

expect :: TokenClass -> Parser ()
expect t = selectToken (\x -> if tokClass x == t then Just () else Nothing)

var :: Parser Text
var = selectToken \case
  Token _ _ _ _ (Tok_var v) -> pure v
  _ -> Nothing

optionally :: Parser a -> Parser (Maybe a)
optionally p = fmap Just p <|> pure Nothing

braces :: Parser a -> Parser a
braces k = do
  expect Tok_obrace
  x <- k
  expect Tok_cbrace
  pure x

parens :: Parser a -> Parser a
parens k = do
  expect Tok_oparen
  x <- k
  expect Tok_cparen
  pure x

instance Alternative Parser where
  empty = Parser \_ -> Left Empty
  Parser kx <|> Parser ky = Parser \x ->
    case kx x of
      Right x -> Right x
      Left e  ->
        case ky x of
          Left _ -> Left e
          Right y -> Right y

attachPos :: Parser RawExpr -> Parser RawExpr
attachPos k = do
  start <- gets (\(ParseState ~(x:_) _ _) -> (tokLine x, tokCol x - (tokOff x - tokSOff x)))
  x <- k
  end <- gets (\(ParseState _ l c) -> (l, c))
  pure (RSrcPos start end x)

body :: Parser RawExpr
body = attachPos letExpr <|> attachPos lamExpr <|> attachPos exprPi where
  letExpr = do
    expect Tok_let
    n <- var
    expect Tok_colon
    t <- body
    letSmol n t <|> letBig n t

  letSmol n t = do
    expect Tok_equal
    d <- body
    expect Tok_semi
    Rlet n t d <$> body

  letBig n t = do
    expect Tok_semi
    selectToken \case
      Token _ _ _ _ (Tok_var n') | n' == n -> Just ()
      _ -> Nothing
    args <- many arg
    expect Tok_equal
    d <- body
    expect Tok_semi
    Rlet n t (foldr lam d args) <$> body

  lamExpr = do
    expect Tok_lambda
    vs <- some arg
    expect Tok_arrow
    e <- body
    pure (foldr lam e vs)

  arg = fmap (Ex,) var <|> fmap (Im,) (braces var)

  lam (p, v) b = Rlam p v b

exprPi :: Parser RawExpr
exprPi = attachPos $
  do
    bs <- optionally binder
    case bs of
      Just k -> foldl (.) id k <$> attachPos exprPi
      Nothing -> attachPos exprArr
  where
    binder = (some (parens (bind Ex) <|> braces (bind Im)) <* expect Tok_arrow)
         <|> (fmap pure (parens sigma) <* expect Tok_times)

    bind p = do
      names <- some var
      expect Tok_colon
      t <- exprPi
      pure (foldr (\n k -> Rpi p n t . k) id names)

    sigma = do
      n <- var
      expect Tok_colon
      Rsigma n <$> exprPi

exprArr :: Parser RawExpr
exprArr = attachPos $ do
  t <- attachPos exprApp
  c <- optionally (fmap (const True) (expect Tok_arrow) <|> fmap (const False) (expect Tok_times))
  case c of
    Just True -> Rpi Ex (T.singleton '_') t <$> exprPi
    Just False -> Rsigma (T.singleton '_') t <$> exprPi
    Nothing -> pure t

exprEq0 :: Parser RawExpr
exprEq0 = attachPos $
  do
    head <- atom
    spine <- many spineEntry
    pure (foldl app head spine)
  where
    spineEntry = fmap (Ex,) atom <|> fmap (Im,) (braces exprPi)
    app f (x, s) = Rapp x f s

exprApp :: Parser RawExpr
exprApp = attachPos $ do
  t <- exprEq0
  c <- optionally (expect Tok_equiv)
  case c of
    Just () -> Req t <$> exprEq0
    Nothing -> pure t

atom0 :: Parser RawExpr
atom0 = attachPos $
       fmap Rvar var
   <|> fmap (const Rtype)  (expect Tok_type)
   <|> fmap (const Rhole)  (expect Tok_under)
   <|> fmap (const Rtop)   (expect Tok_top)
   <|> fmap (const Rrefl)  (expect Tok_refl)
   <|> fmap (const Rcoe)   (expect Tok_coe)
   <|> fmap (const Rcong)  (expect Tok_cong)
   <|> fmap (const Rsym)   (expect Tok_sym)
   <|> fmap (const Runit)  (parens (pure ()))
   <|> parens pair

pair :: Parser RawExpr
pair = attachPos $ do
  t <- body
  c <- optionally (expect Tok_comma)
  case c of
    Just () -> Rpair t <$> pair
    Nothing -> pure t

atom :: Parser RawExpr
atom = attachPos $
  do
    e <- atom0
    c <- many (selectToken (projection . tokClass))
    pure $ case c of
      []  -> e
      sls -> foldl (flip ($)) e sls
  where
    projection Tok_proj1 = pure Rproj1
    projection Tok_proj2 = pure Rproj2
    projection _         = Nothing