{- cabal:
build-depends: base, megaparsec, parser-combinators, haskeline, generics-sop, prettyprinter, sop-core, lens, mono-traversable, optics, pretty-simple
default-extensions: LambdaCase, DeriveAnyClass, DeriveGeneric, TypeApplications, FlexibleContexts, PartialTypeSignatures, GADTs, TypeFamilies, TypeOperators, DuplicateRecordFields, OverloadedLabels
-}
{-# LANGUAGE LambdaCase, DeriveAnyClass, DeriveGeneric, TypeApplications, FlexibleContexts, PartialTypeSignatures, GADTs, TypeFamilies, TypeOperators, DuplicateRecordFields, OverloadedLabels #-}
module Main where

import Data.List ((\\))
import Data.Void
import Data.Function
import Debug.Trace
import Data.Char
import Text.Pretty.Simple

import System.Console.Haskeline
import System.IO
import Control.Monad.IO.Class
import Control.Monad.Combinators.Expr
import Data.Word
import Data.Data
import qualified Text.Megaparsec.Char.Lexer as L
import Text.Megaparsec
import Text.Megaparsec.Char

data Expr =
    App Expr Expr
  | Abs String Expr
  | Var String
  | Lit Literal
  | Let String Expr Expr
  | BinOp String Expr Expr
  deriving (Show, Eq)
  
  
data Literal = LitInt Int | LitStr String | LitB Bool
  deriving (Show, Eq)



type Prec = Word8

opP :: Parser String
opP = lexeme (some (satisfy (flip any [isPunctuation, isSymbol] . (&)))) <?> "symbol"

operators :: MonadFail m => String -> m (Prec, Prec)
operators = \case
  "+"  -> pure (1, 2)
  "-"  -> pure (1, 2)
  "++" -> pure (1, 2)
  "&&" -> pure (1, 2)
  "||" -> pure (1, 2)
  "*"  -> pure (3, 4)
  "/"  -> pure (3, 4)
  _   -> fail "invalid operator"

recexpr :: Prec -> Parser Expr
recexpr min_prec = do
  lhs  <- try (parens $ recexpr 0) <|> valP
  go lhs
  where
    go lhs = do
      withRecovery (\_ -> pure lhs) $ do
        op <- lookAhead opP
        (lbp, rbp) <- operators op
        if lbp < min_prec then
            pure lhs
        else do
            opP
            rhs <- recexpr rbp
            go (BinOp op lhs rhs)

type Parser = Parsec Void String

main :: IO ()
main = parseTest (recexpr 0 <* eof) input >> pPrint (fmap eval $ parseMaybe (recexpr 0) input)

input = "let y = \\x -> x + 1 in y 1"

freeVars :: Expr -> [String]
freeVars = \case
  App one two  -> freeVars one ++ freeVars two
  Abs var body -> freeVars body \\ [var]
  Var var      -> [var]
  _            -> []

rewrite :: String -> Expr -> Expr -> Expr
rewrite var impl replace = case impl of
  App one two            -> App (rewrite var one replace) (rewrite var two replace)
  Abs name body          -> Abs name (rewrite var body replace)
  Var name | name == var -> replace
  BinOp op p0 p1         -> BinOp op (rewrite var p0 replace) (rewrite var p1 replace)
  other                  -> other

eval :: Expr -> Expr
eval = \case
  App (Abs var impl) replace           -> eval   (rewrite var impl replace)
  Let var replace impl                 -> eval   (rewrite var impl replace)
  BinOp "*" p0 p1                      -> opInt  (*)  (eval p0) (eval p1)
  BinOp "/" p0 p1                      -> opInt  div  (eval p0) (eval p1)
  BinOp "+" p0 p1                      -> opInt  (+)  (eval p0) (eval p1)
  BinOp "-" p0 p1                      -> opInt  (-)  (eval p0) (eval p1)
  BinOp "++" p0 p1                     -> opStr  (++) (eval p0) (eval p1)
  BinOp "&&" p0 p1                     -> opBool (&&) (eval p0) (eval p1)
  BinOp "||" p0 p1                     -> opBool (||) (eval p0) (eval p1)
  rest                                 -> rest

eval_ :: Expr -> IO ()
eval_ = pPrint . eval

opInt f (Lit (LitInt x)) (Lit (LitInt y)) = Lit (LitInt (f x y))

opStr f (Lit (LitStr x)) (Lit (LitStr y)) = Lit (LitStr (f x y))

opBool f (Lit (LitB x)) (Lit (LitB y)) = Lit (LitB (f x y))

--appP :: Parser Expr
--appP = do
--  left <- absP
--  rightOpt <- optional $ char ' ' *> recexpr 0
--  case rightOpt of
--    Nothing -> pure left
--    Just right -> pure $ App left right


appP :: Parser Expr
appP = do
  left <- absP
  rightOpt <- lexeme (recexpr 0)
  pure $ App left rightOpt

absP :: Parser Expr
absP = do _ <- symbol "\\"
          vars <- some (lexeme word)
          _ <- symbol "->"
          impl <- recexpr 0
          pure $ foldr Abs impl vars


valP :: Parser Expr
valP = parens (recexpr 0) <|> letP <|> absP <|> appP <|> litP <|> varP

varP :: Parser Expr
varP = Var <$> lexeme word

letP :: Parser Expr
letP = do _ <- symbol "let"
          var <- lexeme word
          _ <- symbol "="
          impl <- lexeme (recexpr 0)
          _ <- symbol "in"
          rest <- lexeme (recexpr 0)
          pure $ Let var impl rest

litP :: Parser Expr
litP = Lit <$> lexeme (intP <|> boolP <|> stringP)

intP :: Parser Literal
intP = LitInt <$> lexeme L.decimal

boolP :: Parser Literal
boolP = LitB <$> (trueL <|> falseL)
  where
    trueL = symbol "true" *> pure True
    falseL = symbol "false" *> pure False

stringP :: Parser Literal
stringP = LitStr <$> stringL

stringL :: Parser String
stringL = string "\"" >> manyTill L.charLiteral (string "\"")

parens :: Parser a -> Parser a
parens = between (symbol "(") (symbol ")")

symbol :: String -> Parser String
symbol = L.symbol space

word :: Parser String
word = some alphaNumChar

lexeme :: Parser a -> Parser a
lexeme = L.lexeme space