module Hint.Lexer where

import           Control.Applicative (Alternative (..), asum)
import           Data.Char           (isAlpha, isDigit, isSpace)

type Input = String

data SourcePos = SourcePos
  { line   :: Int
  , column :: Int}
  deriving (Show, Eq)

data LexerState = LexerState
  { sourcepos :: SourcePos
  , input     :: String
  } deriving (Show, Eq)

data LexError
  = UnexpectedEof
  | IllegalChar Char
  | UnexpectedChar SourcePos [Char]
  | NoMatch
  -- ^ Used for Alternative implementation
  deriving (Show, Eq)

data TokenKind
  = TIdent String
  | TInt Int
  | TString String
  | TOpenParen
  | TCloseParen
  | TOpenCurly
  | TCloseCurly
  | TOpenSquare
  | TCloseSquare
  | TComma
  | TEof
  | TPlus
  | TMinus
  | TEquals
  | KeywordLet
  | KeywordIn
  deriving (Show, Eq)

data Span = Span SourcePos SourcePos
  deriving (Show, Eq)

data Token = Token TokenKind Span
  deriving (Show, Eq)

getKind :: Token -> TokenKind
getKind (Token kind _) = kind

newtype Lexer a = Lexer
  { unLexer :: LexerState -> Either LexError (a, LexerState) }

instance Functor Lexer where
  fmap f (Lexer lf) = Lexer $ \s -> do
    (c, s') <- lf s
    pure (f c, s')

instance Applicative Lexer where
  pure a = Lexer $ \s -> Right (a, s)

  Lexer lf <*> Lexer lx = Lexer $ \s -> do
    (f, s') <- lf s
    (x, s'') <- lx s'
    pure (f x, s'')

instance Alternative Lexer where
  empty = Lexer $ \_ -> Left NoMatch

  Lexer p <|> Lexer q = Lexer $ \s ->
    case p s of
      Right n -> Right n
      Left _  -> q s

instance Monad Lexer where
  return = pure

  Lexer lx >>= f = Lexer $ \s -> do
    (x, s') <- lx s
    unLexer (f x) s'

-- | Toplevel interface to lexer
runLexer :: Lexer a -> Input -> Either LexError (a, LexerState)
runLexer lx = unLexer lx . LexerState (SourcePos 1 1)

runProd :: Input -> Either LexError [Token]
runProd = fmap fst . runLexer gram

-- | Lifts `TokenKind` lexers into Token by tracking span information behind the
-- scenes.
token :: Lexer TokenKind -> Lexer Token
token (Lexer p) = Lexer $ \st -> do
  let start = sourcepos st
  (kind, st') <- p st
  let end = sourcepos st'

  pure $ (Token kind (Span start end), st')

lexChar :: Lexer Char
lexChar = Lexer $ \st ->
  case input st of
    []            -> Left UnexpectedEof
    ('\n' : rest) -> Right ('\n', st {sourcepos = withNewLine st, input = rest})
    (c : rest)    -> Right (c, st { sourcepos = withNewCol st, input = rest })
  where
    withNewLine, withNewCol :: LexerState -> SourcePos

    withNewLine = liftA2 SourcePos ((+ 1) . line . sourcepos) (pure 1)
    withNewCol = liftA2 SourcePos (line . sourcepos) ((+ 1) . column . sourcepos)

-- lexEof :: Lexer TokenKind
-- lexEof = Lexer $ \st ->
--   case input st of
--     [] -> pure (TEof, st)
--     _ -> Left $ UnexpectedChar $ sourcepos st

sat :: (Char -> Bool) -> Lexer Char
sat p = do
  c <- lexChar
  if p c then pure c else Lexer $ \s -> Left (UnexpectedChar (sourcepos s) [c])

string :: String -> Lexer String
string = traverse (\c -> sat (== c))

lexKeyword :: Lexer TokenKind
lexKeyword = asum $
  [ KeywordLet <$ string "let"
  , KeywordIn <$ string "in"
  ]

lexIdent :: Lexer TokenKind
lexIdent = TIdent <$> liftA2 (:) (sat isAlpha) (many rest)
  where
    -- characters allowed ofter the initial a-z
    rest = asum $ sat <$> [isAlpha, isDigit, (== '\'')]

lexNumber :: Lexer TokenKind
lexNumber = TInt . read <$> some (sat isDigit)

lexString :: Lexer TokenKind
lexString = TString <$> (sat (== '\"') *> many (sat (/= '\"')) <* sat (== '\"'))

spaces :: Lexer [()]
spaces = many $ () <$ sat isSpace

gram :: Lexer [Token]
gram = some $ spaces *> asum lexers
  where
    lexers = map token $
      [ TOpenParen <$ sat (== '(')
      , TCloseParen <$ sat (== ')')
      , TOpenCurly <$ sat (== '{')
      , TCloseCurly <$ sat (== '}')
      , TOpenSquare <$ sat (== '[')
      , TCloseSquare <$ sat (== ']')
      , TComma <$ sat (== ',')
      , TPlus <$ sat (== '+')
      , TMinus <$ sat (== '-')
      , TEquals <$ sat (== '=')
      , lexKeyword
      , lexIdent
      , lexNumber
      , lexString
      ]