module LanguageTwo where
import qualified PythonAst as Py
import Data.Char
import Data.Functor
import Text.Parsec
import Text.Parsec.Char
import Text.Parsec.Combinator

{- Data Types -}
data Op
    = Add
    | Subtract
    | Multiply
    | Divide
    | Equal
    | NotEqual
    | And
    | Or

data Expr
    = IntLiteral Int
    | BinOp Op Expr Expr
    | Var String
    | Length Expr

data Stmt
    = IfElse Expr Stmt (Maybe Stmt)
    | Assign String Expr
    | Block [Stmt]

data Prog = Prog Expr [Stmt] [Stmt]

{- Parser -}
type Parser = Parsec String ()

parseKw :: String -> Parser ()
parseKw s = string s $> ()

parseKwIf :: Parser ()
parseKwIf = parseKw "if"

parseKwElse :: Parser ()
parseKwElse = parseKw "else"

parseKwState :: Parser ()
parseKwState = parseKw "state"

parseKwEffect :: Parser ()
parseKwEffect = parseKw "effect"

parseKwCombine :: Parser ()
parseKwCombine = parseKw "combine"

parseOp :: String -> Op -> Parser Op
parseOp s o = string s $> o

parseInt :: Parser Int
parseInt = read <$> (many1 digit <* spaces)

parseVar :: Parser String
parseVar =
    do
        c <- satisfy $ \c -> isLetter c || c == '_'
        cs <- many (satisfy isLetter <|> digit) <* spaces
        let name = c:cs
        if name `elem` ["if", "else", "state", "effect", "combine"]
        then fail "Can't use reserved keyword as identifier"
        else return name

parseSurrounded :: Char -> Char -> Parser a -> Parser a
parseSurrounded c1 c2 pe =
    do
        char c1 >> spaces
        e <- pe
        spaces >> char c2 >> spaces
        return e

parseLength :: Parser Expr
parseLength = Length <$> parseSurrounded '|' '|' parseExpr

parseParenthesized :: Parser Expr
parseParenthesized = parseSurrounded '(' ')' parseExpr

parseBasic :: Parser Expr
parseBasic = choice
    [ IntLiteral <$> parseInt
    , Var <$> parseVar
    , parseLength
    , parseParenthesized
    ]

parseLevel :: Parser Op -> Parser Expr -> Parser Expr
parseLevel po pe =
    do
        e <- pe <* spaces
        ops <- many ((flip . BinOp <$> (po <* spaces) <*> pe) <* spaces)
        return $ foldl (flip ($)) e ops

parseExpr :: Parser Expr
parseExpr = foldl (flip parseLevel) parseBasic
    [ parseOp "*" Multiply <|> parseOp "/" Divide
    , parseOp "+" Add <|> parseOp "-" Subtract
    , parseOp "==" Equal <|> parseOp "!=" NotEqual
    , parseOp "&&" And
    , try $ parseOp "||" Or
    ]

parseIf :: Parser Stmt
parseIf = do
    parseKwIf >> spaces
    c <- parseParenthesized
    t <- parseStmt <* spaces
    e <- (Just <$> (parseKwElse >> spaces *> parseStmt)) <|> return Nothing
    return $ IfElse c t e

parseBlockStmts :: Parser [Stmt]
parseBlockStmts = parseSurrounded '{' '}' (many parseStmt)

parseBlock :: Parser Stmt
parseBlock = Block <$> parseBlockStmts

parseAssign :: Parser Stmt
parseAssign = Assign <$>
    (parseVar <* spaces <* char '=' <* spaces) <*>
    parseExpr <* (char ';' >> spaces)

parseStmt :: Parser Stmt
parseStmt = choice
    [ parseIf
    , parseAssign
    , parseBlock
    ]

parseProgram :: Parser Prog
parseProgram = do
    state <- parseKwState >> spaces *> parseExpr <* char ';' <* spaces
    effect <- parseKwEffect >> spaces *> parseBlockStmts <* spaces
    combined <- parseKwCombine >> spaces *> parseBlockStmts <* spaces
    return $ Prog state effect combined

parse :: String -> String -> Either ParseError Prog
parse = runParser parseProgram ()

{- Translation -}
baseFunction :: Py.PyExpr -> [Py.PyStmt] -> [Py.PyStmt] -> Py.PyStmt
baseFunction s e c = Py.FunctionDef "prog" ["xs"] $
    [Py.IfElse
        (Py.BinOp Py.LessThan
            (Py.FunctionCall (Py.Var "len") [Py.Var "xs"])
            (Py.IntLiteral 2))
        [Py.Return $ Py.Tuple [s, Py.Var "xs"]]
        []
        Nothing
    , Py.Assign (Py.VarPat "leng")
        (Py.BinOp Py.FloorDiv
            (Py.FunctionCall (Py.Var "len") [Py.Var "xs"]) 
            (Py.IntLiteral 2))
    , Py.Assign (Py.VarPat "left")
        (Py.Access
            (Py.Var "xs") 
            [Py.Slice Nothing $ Just (Py.Var "leng")])
    , Py.Assign (Py.VarPat "right")
        (Py.Access
            (Py.Var "xs") 
            [Py.Slice (Just (Py.Var "leng")) Nothing])
    , Py.Assign (Py.TuplePat [Py.VarPat "ls", Py.VarPat "left"])
        (Py.FunctionCall (Py.Var "prog") [Py.Var "left"])
    , Py.Assign (Py.TuplePat [Py.VarPat "rs", Py.VarPat "right"])
        (Py.FunctionCall (Py.Var "prog") [Py.Var "right"])
    , Py.Standalone $
        Py.FunctionCall (Py.Member (Py.Var "left") "reverse") []
    , Py.Standalone $
        Py.FunctionCall (Py.Member (Py.Var "right") "reverse") []
    , Py.Assign (Py.VarPat "state") s
    , Py.Assign (Py.VarPat "source") (Py.IntLiteral 0)
    , Py.Assign (Py.VarPat "total") (Py.ListLiteral [])
    , Py.While
        (Py.BinOp Py.And
            (Py.BinOp Py.NotEqual (Py.Var "left") (Py.ListLiteral []))
            (Py.BinOp Py.NotEqual (Py.Var "right") (Py.ListLiteral []))) $
        [ Py.IfElse
            (Py.BinOp Py.LessThanEq
                (Py.Access (Py.Var "left") [Py.IntLiteral $ -1])
                (Py.Access (Py.Var "right") [Py.IntLiteral $ -1]))
            [ Py.Standalone $
                Py.FunctionCall (Py.Member (Py.Var "total") "append")
                    [Py.FunctionCall (Py.Member (Py.Var "left") "pop") []]
            , Py.Assign (Py.VarPat "source") (Py.IntLiteral 1)
            ]
            [] $
            Just
            [ Py.Standalone $
                Py.FunctionCall (Py.Member (Py.Var "total") "append")
                    [Py.FunctionCall (Py.Member (Py.Var "right") "pop") []]
            , Py.Assign (Py.VarPat "source") (Py.IntLiteral 2)
            ]
        ] ++ e
    ] ++ c ++
    [ Py.Standalone $ Py.FunctionCall (Py.Member (Py.Var "left") "reverse") []
    , Py.Standalone $ Py.FunctionCall (Py.Member (Py.Var "right") "reverse") []
    , Py.Return $ Py.Tuple
        [ Py.Var "state"
        , foldl (Py.BinOp Py.Add) (Py.Var "total") [Py.Var "left", Py.Var "right"]
        ]
    ]

translateExpr :: Expr -> Py.PyExpr
translateExpr (IntLiteral i) = Py.IntLiteral i
translateExpr (BinOp op l r) =
    Py.BinOp (translateOp op) (translateExpr l) (translateExpr r)
translateExpr (Var s)
    | s == "SOURCE" = Py.Var "source"
    | s == "LEFT" = Py.Var "left"
    | s == "RIGHT" = Py.Var "right"
    | s == "STATE" = Py.Var "state"
    | s == "LSTATE" = Py.Var "ls"
    | s == "RSTATE" = Py.Var "rs"
    | s == "L" = Py.IntLiteral 1
    | s == "R" = Py.IntLiteral 2
    | otherwise = Py.Var s
translateExpr (Length e) = Py.FunctionCall (Py.Var "len") [translateExpr e]

translateOp :: Op -> Py.PyBinOp
translateOp Add = Py.Add
translateOp Subtract = Py.Subtract
translateOp Multiply = Py.Multiply
translateOp Divide = Py.Divide
translateOp Equal = Py.Equal
translateOp NotEqual = Py.NotEqual
translateOp And = Py.And
translateOp Or = Py.Or

translateStmt :: Stmt -> [Py.PyStmt]
translateStmt (IfElse c t e) =
    [Py.IfElse (translateExpr c) (translateStmt t) [] (translateStmt <$> e)]
translateStmt (Assign "STATE" e) = [Py.Assign (Py.VarPat "state") (translateExpr e)]
translateStmt (Assign v e) = [Py.Assign (Py.VarPat v) (translateExpr e)]
translateStmt (Block s) = concatMap translateStmt s

translate :: Prog -> [Py.PyStmt]
translate (Prog s e c) =
    [baseFunction (translateExpr s) (concatMap translateStmt e) (concatMap translateStmt c)]