-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Lexing where

import Control.Monad.State.Strict
import Data.Char
import Data.HashSet qualified as HS
import qualified Data.Scientific as Scientific
import Data.String (fromString)
import Data.Text (Text)
import Data.Text          qualified as T
import Data.Void
import Data.Word
import qualified Data.Map.Strict as M

import Text.Megaparsec hiding (Label, State)
import Text.Megaparsec.Char hiding (space, eol)
import qualified Text.Megaparsec.Char as MC
import qualified Text.Megaparsec.Char.Lexer as L
import Text.Megaparsec.Debug

import Err
import PPrint
import Types.Primitives
import Types.Source
import Util (toSnocList)

data ParseCtx = ParseCtx
  { curIndent      :: Int  -- used Reader-style (i.e. ask/local)
  , canBreak       :: Bool -- used Reader-style (i.e. ask/local)
  , prevWhitespace :: Bool -- tracks whether we just consumed whitespace
  , sourceIdCounter :: Int  -- starts at 1 (0 is reserved for the root)
  , curAtomicLexemes :: [SrcId]
  , curLexemeInfo    :: LexemeInfo } -- append to, writer-style

initParseCtx :: ParseCtx
initParseCtx = ParseCtx 0 False False 1 mempty mempty

type Parser = StateT ParseCtx (Parsec Void Text)

parseit :: Text -> Parser a -> Except a
parseit s p = case parse (fst <$> runStateT p initParseCtx) "" s of
  Left e  -> throwErr $ ParseErr $ MiscParseErr $ errorBundlePretty e
  Right x -> return x

mustParseit :: Text -> Parser a -> a
mustParseit s p  = case parseit s p of
  Success x -> x
  Failure e -> error $ "This shouldn't happen:\n" ++ pprint e

debug :: (Show a) => String -> Parser a -> Parser a
debug lbl action = do
  ctx <- get
  lift $ dbg lbl $ fst <$> runStateT action ctx

-- === Lexemes ===

type Lexer = Parser

nextChar :: Lexer Char
nextChar = do
  i <- getInput
  guard $ not $ T.null i
  return $ T.head i
{-# INLINE nextChar #-}

anyCaseName  :: Lexer (WithSrc SourceName)
anyCaseName = label "name" $ lexeme LowerName anyCaseName' -- TODO: distinguish lowercase/uppercase

anyCaseName' :: Lexer SourceName
anyCaseName' =
  liftM MkSourceName $ checkNotKeyword $ (:) <$> satisfy (\c -> isLower c || isUpper c) <*>
    (T.unpack <$> takeWhileP Nothing (\c -> isAlphaNum c || c == '\'' || c == '_'))

anyName :: Lexer (WithSrc SourceName)
anyName = anyCaseName <|> symName

checkNotKeyword :: Parser String -> Parser String
checkNotKeyword p = try $ do
  s <- p
  failIf (s `HS.member` keyWordSet) $ show s ++ " is a reserved word"
  return s
{-# INLINE checkNotKeyword #-}

data KeyWord = DefKW | ForKW | For_KW | RofKW | Rof_KW | CaseKW | OfKW
             | DataKW | StructKW | InterfaceKW
             | InstanceKW | GivenKW | WithKW | SatisfyingKW
             | IfKW | ThenKW | ElseKW | DoKW
             | ImportKW | ForeignKW | NamedInstanceKW
             | EffectKW | HandlerKW | JmpKW | CtlKW | ReturnKW | ResumeKW
             | CustomLinearizationKW | CustomLinearizationSymbolicKW | PassKW
  deriving (Enum)

keyWordToken :: KeyWord -> String
keyWordToken = \case
  DefKW           -> "def"
  ForKW           -> "for"
  RofKW           -> "rof"
  For_KW          -> "for_"
  Rof_KW          -> "rof_"
  CaseKW          -> "case"
  IfKW            -> "if"
  ThenKW          -> "then"
  ElseKW          -> "else"
  OfKW            -> "of"
  DataKW          -> "enum"
  StructKW        -> "struct"
  InterfaceKW     -> "interface"
  InstanceKW      -> "instance"
  NamedInstanceKW -> "named-instance"
  GivenKW         -> "given"
  WithKW          -> "with"
  SatisfyingKW    -> "satisfying"
  DoKW            -> "do"
  ImportKW        -> "import"
  ForeignKW       -> "foreign"
  EffectKW        -> "effect"
  HandlerKW       -> "handler"
  JmpKW           -> "jmp"
  CtlKW           -> "ctl"
  ReturnKW        -> "return"
  ResumeKW        -> "resume"
  CustomLinearizationKW -> "custom-linearization"
  CustomLinearizationSymbolicKW -> "custom-linearization-symbolic"
  PassKW          -> "pass"

keyWord :: KeyWord -> Lexer ()
keyWord kw = atomicLexeme Keyword $ try $
  string (fromString $ keyWordToken kw) >> notFollowedBy nameTailChar
  where
    nameTailChar :: Parser Char
    nameTailChar = alphaNumChar <|> char '\'' <|> char '_'

keyWordSet :: HS.HashSet String
keyWordSet = HS.fromList keyWordStrs

keyWordStrs :: [String]
keyWordStrs = map keyWordToken [DefKW .. PassKW]

primName :: Lexer (WithSrc String)
primName = lexeme MiscLexeme $ try $ char '%' >> some alphaNumChar

charLit :: Lexer (WithSrc Char)
charLit = lexeme MiscLexeme $ char '\'' >> L.charLiteral <* char '\''

strLit :: Lexer (WithSrc String)
strLit = lexeme StringLiteralLexeme $ char '"' >> manyTill L.charLiteral (char '"')

natLit :: Lexer (WithSrc Word64)
natLit = lexeme LiteralLexeme $ try $ L.decimal <* notFollowedBy (char '.')

doubleLit :: Lexer (WithSrc Double)
doubleLit = lexeme LiteralLexeme $
      try L.float
  <|> try (fromIntegral <$> (L.decimal :: Parser Int) <* char '.')
  <|> try do
    s <- L.scientific
    case Scientific.toBoundedRealFloat s of
      Right f -> return f
      Left  _ -> fail "Non-representable floating point literal"

knownSymStrs :: HS.HashSet String
knownSymStrs = HS.fromList
  [ ".", ":", "::", "!", "=", "-", "+", "||", "&&"
  , "$", "&>", "|", ",", ",>", "<-", "+=", ":="
  , "->", "->>", "=>", "?->", "?=>", "<<<", ">>>"
  , "..", "<..", "..<", "..<", "<..<", "?", "#", "##", "#?", "#&", "#|", "@"]

sym :: Text -> Lexer ()
sym s = atomicLexeme Symbol $ sym' s

symWithId :: Text -> Lexer SrcId
symWithId s = liftM srcPos $ lexeme Symbol $ sym' s

-- string must be in `knownSymStrs`
sym' :: Text -> Lexer ()
sym' s = void $ try $ string s >> notFollowedBy symChar

anySym :: Lexer (WithSrc String)
anySym = lexeme Symbol $ try $ do
  s <- some symChar
  failIf (s `HS.member` knownSymStrs) ""
  return s

symName :: Lexer (WithSrc SourceName)
symName = label "symbol name" $ lexeme Symbol $ try $ do
  s <- between (char '(') (char ')') $ some symChar
  return $ MkSourceName $ "(" <> s <> ")"

backquoteName :: Lexer (WithSrc SourceName)
backquoteName = label "backquoted name" $
  lexeme Symbol $ try $ between (char '`') (char '`') anyCaseName'

-- brackets and punctuation
-- (can't treat as sym because e.g. `((` is two separate lexemes)
lParen, rParen, lBracket, rBracket, lBrace, rBrace, semicolon, underscore :: Lexer ()

lParen    = charLexeme '('
rParen    = charLexeme ')'
lBracket  = charLexeme '['
rBracket  = charLexeme ']'
lBrace    = charLexeme '{'
rBrace    = charLexeme '}'
semicolon = charLexeme ';'
underscore = charLexeme '_'

charLexeme :: Char -> Parser ()
charLexeme c = atomicLexeme Symbol $ void $ char c

symChar :: Parser Char
symChar = token (\c -> if HS.member c symChars then Just c else Nothing) mempty

symChars :: HS.HashSet Char
symChars = HS.fromList ".,!$^&*:-~+/=<>|?\\@#"

-- XXX: unlike other lexemes, this doesn't consume trailing whitespace
dot :: Parser ()
dot = do
  WithSrc sid () <- lexeme' (return ()) Symbol (void $ char '.')
  emitAtomicLexeme sid

-- === Util ===

sc :: Parser ()
sc = (skipSome s >> recordWhitespace) <|> return ()
  where s = hidden space <|> hidden lineComment

lineComment :: Parser ()
lineComment = string "#" >> void (takeWhileP (Just "char") (/= '\n'))

outputLines :: Parser ()
outputLines = void $ many (symbol ">" >> takeWhileP Nothing (/= '\n') >> ((eol >> return ()) <|> eof))

space :: Parser ()
space = gets canBreak >>= \case
  True  -> space1
  False -> void $ takeWhile1P (Just "white space") (`elem` (" \t" :: String))

setCanBreakLocally :: Bool -> Parser a -> Parser a
setCanBreakLocally brLocal p = do
  brPrev <- gets canBreak
  modify \ctx -> ctx {canBreak = brLocal}
  ans <- p
  modify \ctx -> ctx {canBreak = brPrev}
  return ans
{-# INLINE setCanBreakLocally #-}

mayBreak :: Parser a -> Parser a
mayBreak p = setCanBreakLocally True p
{-# INLINE mayBreak #-}

mayNotBreak :: Parser a -> Parser a
mayNotBreak p = setCanBreakLocally False p
{-# INLINE mayNotBreak #-}

precededByWhitespace :: Parser Bool
precededByWhitespace = gets prevWhitespace
{-# INLINE precededByWhitespace #-}

recordWhitespace :: Parser ()
recordWhitespace = modify \ctx -> ctx { prevWhitespace = True }
{-# INLINE recordWhitespace #-}

recordNonWhitespace :: Parser ()
recordNonWhitespace = modify \ctx -> ctx { prevWhitespace = False }
{-# INLINE recordNonWhitespace #-}

nameString :: Parser String
nameString = lexemeIgnoreSrcId LowerName . try $ (:) <$> lowerChar <*> many alphaNumChar

thisNameString :: Text -> Parser ()
thisNameString s = lexemeIgnoreSrcId MiscLexeme $ try $ string s >> notFollowedBy alphaNumChar

bracketed :: Parser () -> Parser () -> Parser a -> Parser a
bracketed left right p = do
  left
  ans <- mayBreak $ sc >> p
  right
  return ans
{-# INLINE bracketed #-}

braces :: Parser a -> Parser a
braces p = bracketed lBrace rBrace p
{-# INLINE braces #-}

nextLine :: Parser ()
nextLine = do
  eol
  n <- curIndent <$> get
  void $ mayNotBreak $ many $ try (sc >> eol)
  void $ replicateM n (char ' ')

withSource :: Parser a -> Parser (Text, a)
withSource p = do
  s <- getInput
  start <- getOffset
  x <- p
  end <- getOffset
  return (T.take (end - start) s, x)
{-# INLINE withSource #-}

withIndent :: Parser a -> Parser a
withIndent p = do
  nextLine
  indent <- T.length <$> takeWhileP (Just "space") (==' ')
  when (indent <= 0) empty
  locallyExtendCurIndent indent $ mayNotBreak p
{-# INLINE withIndent #-}

locallyExtendCurIndent :: Int -> Parser a -> Parser a
locallyExtendCurIndent n p = do
  indentPrev <- gets curIndent
  modify \ctx -> ctx { curIndent = indentPrev + n }
  ans <- p
  modify \ctx -> ctx { curIndent = indentPrev }
  return ans

eol :: Parser ()
eol = void MC.eol

eolf :: Parser ()
eolf = eol <|> eof

failIf :: Bool -> String -> Parser ()
failIf True s = fail s
failIf False _ = return ()

freshSrcId :: Parser SrcId
freshSrcId = do
  c <- gets sourceIdCounter
  modify \ctx -> ctx { sourceIdCounter = c + 1 }
  return $ SrcId c

withLexemeInfo :: Parser a -> Parser (LexemeInfo, a)
withLexemeInfo cont = do
  smPrev <- gets curLexemeInfo
  modify \ctx -> ctx { curLexemeInfo = mempty }
  result <- cont
  sm <- gets curLexemeInfo
  modify \ctx -> ctx { curLexemeInfo = smPrev }
  return (sm, result)

emitLexemeInfo :: LexemeInfo -> Parser ()
emitLexemeInfo m = modify \ctx -> ctx { curLexemeInfo = curLexemeInfo ctx <> m }

lexemeIgnoreSrcId :: LexemeType -> Parser a -> Parser a
lexemeIgnoreSrcId lexemeType p = withoutSrc <$> lexeme lexemeType p

symbol :: Text -> Parser ()
symbol s = void $ L.symbol sc s

lexeme :: LexemeType -> Parser a -> Parser (WithSrc a)
lexeme lexemeType p = lexeme' sc lexemeType p
{-# INLINE lexeme #-}

lexeme' :: Parser () -> LexemeType -> Parser a -> Parser (WithSrc a)
lexeme' sc' lexemeType p = do
  start <- getOffset
  ans <- p
  end <- getOffset
  recordNonWhitespace
  sc'
  sid <- freshSrcId
  emitLexemeInfo $ mempty
    { lexemeList = toSnocList [sid]
    , lexemeInfo = M.singleton sid (lexemeType, (start, end)) }
  return $ WithSrc sid ans
{-# INLINE lexeme' #-}

atomicLexeme :: LexemeType -> Parser () -> Parser ()
atomicLexeme lexemeType p = do
  WithSrc sid () <- lexeme lexemeType p
  emitAtomicLexeme sid
{-# INLINE atomicLexeme #-}

emitAtomicLexeme :: LexemeId -> Parser ()
emitAtomicLexeme sid = modify \ctx ->
  ctx { curAtomicLexemes = curAtomicLexemes ctx ++ [sid] }

collectAtomicLexemeIds :: Parser a -> Parser ([SrcId], a)
collectAtomicLexemeIds p = do
  prevAtomicLexemes <- gets curAtomicLexemes
  modify \ctx -> ctx { curAtomicLexemes = [] }
  ans <- p
  localLexemes <- gets curAtomicLexemes
  modify \ctx -> ctx { curAtomicLexemes = prevAtomicLexemes }
  return (localLexemes, ans)