{-# 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


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))

eof :: Parser ()
eof = Parser $ \state ->
  case ptTs state of
    [] -> Right ((), state)
    (x:_) -> Left (Unexpected x)

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

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

square :: Parser a -> Parser a
square k = do
  expect Tok_osquare
  x <- k
  expect Tok_csquare
  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 Exp -> Parser Exp
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 (Span start end x)

body :: Parser Exp
body = attachPos letExpr <|> attachPos lamExpr <|> attachPos exprPi where
  lamExpr = do
    expect Tok_lambda
    vs <- some arg
    expect Tok_arrow
    e <- body
    pure (foldr Lam e vs)

  letExpr = do
    expect Tok_let
    v <- T.unpack <$> var
    expect Tok_colon
    t <- body
    expect Tok_equal
    b <- body
    expect Tok_in
    Let v t b <$> body

  arg = T.unpack <$> var

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

    bind = do
      names <- some (T.unpack <$> var)
      expect Tok_colon
      t <- exprPi
      pure (foldr (\n k -> Pi n t . k) id names)

    sigma = do
      names <- some (T.unpack <$> var)
      expect Tok_colon
      t <- exprPi
      pure (foldr (\n k -> Sigma n t . k) id names)

exprArr :: Parser Exp
exprArr = attachPos $ do
  t <- attachPos exprConj
  c <- optionally (fmap (const True) (expect Tok_arrow) <|> fmap (const False) (expect Tok_times))
  case c of
    Just True -> Pi "_" t <$> exprPi
    Just False -> Sigma "_" t <$> exprPi
    Nothing -> pure t

exprApp :: Parser Exp
exprApp = attachPos $
  do
    head <- atom
    spine <- many spineEntry
    pure (foldl app head spine)
  where
    spineEntry = atom
    app f s = App f s

exprDisj :: Parser Exp
exprDisj = attachPos $
  do
    first <- exprApp
    rest <- many disjunct
    pure (foldl IOr first rest)
  where
    disjunct = expect Tok_or *> exprApp

exprConj :: Parser Exp
exprConj = attachPos $
  do
    first <- exprDisj
    rest <- many conjunct
    pure (foldl IAnd first rest)
  where
    conjunct = expect Tok_and *> exprDisj

atom0 :: Parser Exp
atom0 = attachPos $
       fmap (Var . T.unpack) var
   <|> fmap (const Type)     (expect Tok_type)
   <|> fmap (const I)        (expect Tok_I)
   <|> fmap (const I0)       (expect Tok_I0)
   <|> fmap (const I1)       (expect Tok_I1)
   <|> fmap (const Path)     (expect Tok_path)
   <|> fmap (const SubT)     (expect Tok_sub)
   <|> fmap (const PartialT) (expect Tok_phi)
   <|> fmap (const Comp)     (expect Tok_comp)
   <|> fmap INot           (expect Tok_not *> atom)
   <|> parens pair
   <|> square (Partial <$> (system <|> pure []))

atom :: Parser Exp
atom = attachPos $
  do
    e <- atom0
    c <- many (selectToken (projection . tokClass))
    pure $ case c of
      []  -> e
      sls -> foldl (flip ($)) e sls
  where
    projection Tok_p1 = pure Proj1
    projection Tok_p2 = pure Proj2
    projection _      = Nothing


system :: Parser [(Formula, Exp)]
system =
  do
    t <- comp
    x <- optionally (expect Tok_comma)
    case x of
      Just () -> (t:) <$> system
      Nothing -> pure [t]
  where
    comp = do
      t <- formula
      expect Tok_arrow
      (t,) <$> body

pair :: Parser Exp
pair = do
  t <- body
  x <- optionally (expect Tok_comma)
  case x of
    Just () -> Pair t <$> pair
    Nothing -> pure t

statement :: Parser Statement
statement = (assume <|> declare <|> (Eval <$> body)) <* eof where
  assume = do
    expect Tok_assume
    Assume <$> vars

  declare = do
    expect Tok_let
    x <- T.unpack <$> var
    expect Tok_colon
    ty <- body
    expect Tok_equal
    Declare x ty <$> body

  bind = do
    var <- some (T.unpack <$> var)
    expect Tok_colon
    body <- body
    pure $ map ((, body)) var

  vars = do
    var <- bind
    t <- optionally (expect Tok_comma)
    case t of
      Nothing -> pure var
      Just x -> (var ++) <$> vars

formula :: Parser Formula
formula = conjunction where
  conjunction, disjunction, atom :: Parser Formula
  conjunction = do
    d <- disjunction
    t <- optionally (expect Tok_and)
    case t of
      Nothing -> pure d
      Just x -> And d <$> conjunction

  disjunction = do
    d <- atom
    t <- optionally (expect Tok_or)
    case t of
      Nothing -> pure d
      Just x -> Or d <$> disjunction

  atom = (Is1 . T.unpack) <$> var
     <|> (Is0 . T.unpack) <$> (expect Tok_not *> var)
     <|> Top <$ expect Tok_I1
     <|> Bot <$ expect Tok_I0