module Tessera.Parser (parseProgram, parseExpression, tokenize) where

import Data.Char (isAlpha, isAlphaNum, isDigit, isLower, isUpper)
import Data.List (foldl')
import qualified Data.Map.Strict as Map

import Tessera.Syntax

data Token
  = TkLower String
  | TkUpper String
  | TkInt Integer
  | TkSym String
  | TkEof
  deriving (Eq, Show)

data Located = Located
  { locatedSpan :: Span
  , locatedToken :: Token
  } deriving (Eq, Show)

data PState = PState
  { psTokens :: [Located]
  , psArities :: Map.Map String Int
  , psTypeArities :: Map.Map String Int
  }

newtype Parser a = Parser { runParser :: PState -> Either Diagnostic (a, PState) }

instance Functor Parser where
  fmap f (Parser g) = Parser $ \s -> case g s of
    Left e -> Left e
    Right (a, s') -> Right (f a, s')

instance Applicative Parser where
  pure a = Parser $ \s -> Right (a, s)
  Parser f <*> Parser g = Parser $ \s -> case f s of
    Left e -> Left e
    Right (h, s') -> case g s' of
      Left e -> Left e
      Right (a, s'') -> Right (h a, s'')

instance Monad Parser where
  return = pure
  Parser g >>= f = Parser $ \s -> case g s of
    Left e -> Left e
    Right (a, s') -> runParser (f a) s'

parseError :: Span -> String -> Parser a
parseError sp message = Parser $ \_ -> Left (Diagnostic sp message)

peekToken :: Parser Located
peekToken = Parser $ \s -> case psTokens s of
  [] -> Right (Located (Span 0 0) TkEof, s)
  (t : _) -> Right (t, s)

advanceToken :: Parser Located
advanceToken = Parser $ \s -> case psTokens s of
  [] -> Right (Located (Span 0 0) TkEof, s)
  (t : ts) -> Right (t, s { psTokens = ts })

withArities :: (Map.Map String Int -> Map.Map String Int) -> Parser ()
withArities f = Parser $ \s -> Right ((), s { psArities = f (psArities s) })

lookupArity :: Span -> String -> Parser Int
lookupArity sp name = Parser $ \s -> case Map.lookup name (psArities s) of
  Just n -> Right (n, s)
  Nothing -> Left (Diagnostic sp ("unknown constructor " ++ name))

builtinTypeArities :: Map.Map String Int
builtinTypeArities = Map.fromList [("Int", 0), ("Bool", 0)]

withTypeArities :: (Map.Map String Int -> Map.Map String Int) -> Parser ()
withTypeArities f = Parser $ \s -> Right ((), s { psTypeArities = f (psTypeArities s) })

lookupTypeArity :: Span -> String -> Parser Int
lookupTypeArity sp name = Parser $ \s -> case Map.lookup name (psTypeArities s) of
  Just n -> Right (n, s)
  Nothing -> Left (Diagnostic sp ("unknown type " ++ name))

matchSymbol :: String -> Parser ()
matchSymbol symbol = do
  located <- peekToken
  case locatedToken located of
    TkSym s | s == symbol -> advanceToken >> return ()
    _ -> parseError (locatedSpan located) ("expected " ++ symbol)

matchLower :: Parser String
matchLower = do
  located <- peekToken
  case locatedToken located of
    TkLower name -> advanceToken >> return name
    _ -> parseError (locatedSpan located) "expected a lowercase name"

matchUpper :: Parser String
matchUpper = do
  located <- peekToken
  case locatedToken located of
    TkUpper name -> advanceToken >> return name
    _ -> parseError (locatedSpan located) "expected an uppercase constructor name"

matchInt :: Parser Integer
matchInt = do
  located <- peekToken
  case locatedToken located of
    TkInt n -> advanceToken >> return n
    _ -> parseError (locatedSpan located) "expected an integer literal"

matchEof :: Parser ()
matchEof = do
  located <- peekToken
  case locatedToken located of
    TkEof -> return ()
    _ -> parseError (locatedSpan located) "unexpected trailing input"

currentSpan :: Parser Span
currentSpan = locatedSpan <$> peekToken

isSymbol :: String -> Located -> Bool
isSymbol symbol located = locatedToken located == TkSym symbol

parseProgram :: String -> Either Diagnostic Program
parseProgram input = case tokenize input of
  Left e -> Left e
  Right tokens -> case runParser program (PState tokens Map.empty builtinTypeArities) of
    Left e -> Left e
    Right (declarations, _) -> Right declarations

program :: Parser Program
program = do
  located <- peekToken
  case locatedToken located of
    TkEof -> return []
    _ -> do
      declaration <- parseDeclaration
      rest <- program
      return (declaration : rest)

parseDeclaration :: Parser Declaration
parseDeclaration = do
  located <- peekToken
  case locatedToken located of
    TkLower "data" -> parseData
    TkLower "match" -> parseMatch
    _ -> parseError (locatedSpan located) "expected data or match"

parseData :: Parser Declaration
parseData = do
  matchLower
  name <- matchUpper
  parameters <- parseParameters
  withTypeArities (Map.insert name (length parameters))
  matchSymbol "="
  constructors <- parseConstructors
  withArities $ \m ->
    foldl' (\acc c -> Map.insert (constructorName c) (length (constructorArgs c)) acc) m constructors
  return (DataDeclaration name parameters constructors)

parseParameters :: Parser [String]
parseParameters = do
  located <- peekToken
  case locatedToken located of
    TkLower name -> do
      advanceToken
      rest <- parseParameters
      return (name : rest)
    _ -> return []

parseConstructors :: Parser [Constructor]
parseConstructors = do
  name <- matchUpper
  arguments <- parseTypeAtoms
  located <- peekToken
  if isSymbol "|" located
    then do
      advanceToken
      rest <- parseConstructors
      return (Constructor name arguments : rest)
    else return [Constructor name arguments]

parseMatch :: Parser Declaration
parseMatch = do
  matchLower
  name <- matchLower
  arguments <- parseMatchArguments
  matchSymbol ":"
  result <- parseType
  matchSymbol "="
  clauses <- parseClauses
  return (MatchDeclaration name arguments result clauses)

parseMatchArguments :: Parser [(String, Type)]
parseMatchArguments = do
  located <- peekToken
  if isSymbol "(" located
    then do
      advanceToken
      argument <- matchLower
      matchSymbol ":"
      ty <- parseType
      matchSymbol ")"
      rest <- parseMatchArguments
      return ((argument, ty) : rest)
    else return []

parseClauses :: Parser [Clause]
parseClauses = do
  located <- peekToken
  if isSymbol "|" located
    then do
      span' <- currentSpan
      advanceToken
      patterns <- parseClausePatterns
      matchSymbol "->"
      body <- expression
      rest <- parseClauses
      return (Clause span' patterns body : rest)
    else return []

parseClausePatterns :: Parser [Pattern]
parseClausePatterns = do
  located <- peekToken
  if isSymbol "->" located || isSymbol "|" located
    then return []
    else do
      pattern <- parsePattern
      rest <- parseClausePatterns
      return (pattern : rest)

isReserved :: String -> Bool
isReserved name = name `elem` ["data", "match", "if", "then", "else", "let", "in", "true", "false"]

parseType :: Parser Type
parseType = parseTypeAtom

parseTypeArguments :: Int -> Parser [Type]
parseTypeArguments 0 = return []
parseTypeArguments n = do
  first <- parseTypeAtom
  rest <- parseTypeArguments (n - 1)
  return (first : rest)

parseTypeAtom :: Parser Type
parseTypeAtom = do
  located <- peekToken
  case locatedToken located of
    TkLower name | not (isReserved name) -> advanceToken >> return (TVar name)
    TkUpper name -> do
      arity <- lookupTypeArity (locatedSpan located) name
      advanceToken
      arguments <- parseTypeArguments arity
      return (TName name arguments)
    TkSym "(" -> do
      advanceToken
      first <- parseType
      located' <- peekToken
      if isSymbol "," located'
        then do
          advanceToken
          rest <- parseTupleTypes
          matchSymbol ")"
          return (TTuple (first : rest))
        else do
          matchSymbol ")"
          return first
    _ -> parseError (locatedSpan located) "expected a type"

parseTupleTypes :: Parser [Type]
parseTupleTypes = do
  ty <- parseType
  located <- peekToken
  if isSymbol "," located
    then do
      advanceToken
      rest <- parseTupleTypes
      return (ty : rest)
    else return [ty]

parseTypeAtoms :: Parser [Type]
parseTypeAtoms = do
  located <- peekToken
  case locatedToken located of
    TkLower name | not (isReserved name) -> advanceToken >> fmap (TVar name :) parseTypeAtoms
    TkUpper name -> do
      arity <- lookupTypeArity (locatedSpan located) name
      advanceToken
      arguments <- parseTypeArguments arity
      fmap (TName name arguments :) parseTypeAtoms
    TkSym "(" -> do
      advanceToken
      first <- parseType
      located' <- peekToken
      if isSymbol "," located'
        then do
          advanceToken
          rest <- parseTupleTypes
          matchSymbol ")"
          fmap (TTuple (first : rest) :) parseTypeAtoms
        else do
          matchSymbol ")"
          fmap (first :) parseTypeAtoms
    _ -> return []

parsePattern :: Parser Pattern
parsePattern = do
  located <- peekToken
  case locatedToken located of
    TkLower "_" -> advanceToken >> return PWild
    TkLower name -> advanceToken >> return (PVar name)
    TkInt n -> advanceToken >> return (PInt n)
    TkUpper "True" -> advanceToken >> return (PBool True)
    TkUpper "False" -> advanceToken >> return (PBool False)
    TkUpper name -> do
      arity <- lookupArity (locatedSpan located) name
      advanceToken
      arguments <- parsePatternArguments arity
      return (PCon name arguments)
    TkSym "(" -> do
      advanceToken
      first <- parsePattern
      located' <- peekToken
      if isSymbol "," located'
        then do
          advanceToken
          rest <- parsePatternTuple
          matchSymbol ")"
          return (PTuple (first : rest))
        else do
          matchSymbol ")"
          return first
    _ -> parseError (locatedSpan located) "expected a pattern"

parsePatternArguments :: Int -> Parser [Pattern]
parsePatternArguments 0 = return []
parsePatternArguments n = do
  pattern <- parsePattern
  rest <- parsePatternArguments (n - 1)
  return (pattern : rest)

parsePatternTuple :: Parser [Pattern]
parsePatternTuple = do
  pattern <- parsePattern
  located <- peekToken
  if isSymbol "," located
    then do
      advanceToken
      rest <- parsePatternTuple
      return (pattern : rest)
    else return [pattern]

parseExpression :: [(String, Int)] -> String -> Either Diagnostic Expr
parseExpression arities input = case tokenize input of
  Left e -> Left e
  Right tokens -> case runParser expression (PState tokens (Map.fromList arities) builtinTypeArities) of
    Left e -> Left e
    Right (e, _) -> Right e

expression :: Parser Expr
expression = do
  located <- peekToken
  case locatedToken located of
    TkLower "if" -> parseIf
    _ -> parseOr

parseIf :: Parser Expr
parseIf = do
  matchLower
  condition <- expression
  matchLower
  yes <- expression
  matchLower
  no <- expression
  return (EIf condition yes no)

binaryLevel :: Parser Expr -> [(String, Op)] -> Parser Expr -> Parser Expr
binaryLevel next operators parseOperand = do
  first <- parseOperand
  loop first
  where
    loop left = do
      located <- peekToken
      case locatedToken located of
        TkSym symbol
          | Just op <- lookup symbol operators -> do
              advanceToken
              right <- parseOperand
              loop (EBin op left right)
        _ -> return left

parseOr :: Parser Expr
parseOr = do
  first <- parseAnd
  loop first
  where
    loop left = do
      located <- peekToken
      if isSymbol "||" located
        then do
          advanceToken
          right <- parseAnd
          loop (EBin OpOr left right)
        else return left

parseAnd :: Parser Expr
parseAnd = do
  first <- parseComparison
  loop first
  where
    loop left = do
      located <- peekToken
      if isSymbol "&&" located
        then do
          advanceToken
          right <- parseComparison
          loop (EBin OpAnd left right)
        else return left

parseComparison :: Parser Expr
parseComparison = do
  left <- parseAdditive
  located <- peekToken
  case locatedToken located of
    TkSym "==" -> advanceToken >> fmap (EBin OpEq left) parseAdditive
    TkSym "!=" -> advanceToken >> fmap (EBin OpNe left) parseAdditive
    TkSym "<" -> advanceToken >> fmap (EBin OpLt left) parseAdditive
    TkSym "<=" -> advanceToken >> fmap (EBin OpLe left) parseAdditive
    TkSym ">" -> advanceToken >> fmap (EBin OpGt left) parseAdditive
    TkSym ">=" -> advanceToken >> fmap (EBin OpGe left) parseAdditive
    _ -> return left

parseAdditive :: Parser Expr
parseAdditive = binaryLevel parseMultiplicative [("+", OpAdd), ("-", OpSub)] parseMultiplicative

parseMultiplicative :: Parser Expr
parseMultiplicative = binaryLevel parseAtom [("*", OpMul), ("/", OpDiv), ("%", OpMod)] parseAtom

parseAtom :: Parser Expr
parseAtom = do
  located <- peekToken
  case locatedToken located of
    TkLower name -> advanceToken >> return (EVar name)
    TkInt n -> advanceToken >> return (EInt n)
    TkUpper "True" -> advanceToken >> return (EBool True)
    TkUpper "False" -> advanceToken >> return (EBool False)
    TkUpper name -> do
      arity <- lookupArity (locatedSpan located) name
      advanceToken
      arguments <- parseExpressionArguments arity
      return (ECon name arguments)
    TkSym "(" -> do
      advanceToken
      first <- expression
      located' <- peekToken
      if isSymbol "," located'
        then do
          advanceToken
          rest <- parseExpressionTuple
          matchSymbol ")"
          return (ETuple (first : rest))
        else do
          matchSymbol ")"
          return first
    _ -> parseError (locatedSpan located) "expected an expression"

parseExpressionArguments :: Int -> Parser [Expr]
parseExpressionArguments 0 = return []
parseExpressionArguments n = do
  argument <- parseAtom
  rest <- parseExpressionArguments (n - 1)
  return (argument : rest)

parseExpressionTuple :: Parser [Expr]
parseExpressionTuple = do
  e <- expression
  located <- peekToken
  if isSymbol "," located
    then do
      advanceToken
      rest <- parseExpressionTuple
      return (e : rest)
    else return [e]

tokenize :: String -> Either Diagnostic [Located]
tokenize = go 1 1
  where
    go line column input = case input of
      [] -> Right [Located (Span line column) TkEof]
      (c : rest)
        | c == '\n' -> go (line + 1) 1 rest
        | c == ' ' || c == '\t' || c == '\r' -> go line (column + 1) rest
        | c == '-' && not (null rest) && head rest == '-' ->
            let (comment, after) = break (== '\n') rest
            in go line (column + length comment) after
        | isAlpha c || c == '_' ->
            let (word, after) = span (\x -> isAlphaNum x || x == '_' || x == '\'') (c : rest)
            in fmap (Located (Span line column) (classifyWord word) :) (go line (column + length word) after)
        | isDigit c ->
            let (digits, after) = span isDigit (c : rest)
            in fmap (Located (Span line column) (TkInt (read digits)) :) (go line (column + length digits) after)
        | otherwise ->
            case matchSymbolAt input of
              Just (symbol, after) ->
                fmap (Located (Span line column) (TkSym symbol) :) (go line (column + length symbol) after)
              Nothing -> Left (Diagnostic (Span line column) ("unexpected character " ++ [c]))

    classifyWord word
      | isUpper (head word) = TkUpper word
      | otherwise = TkLower word

    matchSymbolAt input = firstJust
      [ if take (length symbol) input == symbol then Just (symbol, drop (length symbol) input) else Nothing
      | symbol <- ["->", "==", "!=", "<=", ">=", "&&", "||", "|", "(", ")", "[", "]", ",", ":", "=", "+", "-", "*", "/", "%", "<", ">", ";"]
      ]

    firstJust [] = Nothing
    firstJust (Just x : _) = Just x
    firstJust (Nothing : xs) = firstJust xs