module LanguageTwo where
import qualified PythonAst as Py
import qualified CommonParsing as P
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 ()

parseVar :: Parser String
parseVar = P.var [ "if", "else", "state", "effect", "combine" ]

parseLength :: Parser Expr
parseLength = Length <$> P.surround '|' '|' parseExpr

parseParenthesized :: Parser Expr
parseParenthesized = P.surround '(' ')' parseExpr

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


parseExpr :: Parser Expr
parseExpr = P.precedence BinOp parseBasic
    [ P.op "*" Multiply <|> P.op "/" Divide
    , P.op "+" Add <|> P.op "-" Subtract
    , P.op "==" Equal <|> P.op "!=" NotEqual
    , P.op "&&" And
    , try $ P.op "||" Or
    ]

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

parseBlockStmts :: Parser [Stmt]
parseBlockStmts = P.surround '{' '}' (many parseStmt)

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

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

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

parseProgram :: Parser Prog
parseProgram = do
    state <- P.kwState >> spaces *> parseExpr <* char ';' <* spaces
    effect <- P.kwEffect >> spaces *> parseBlockStmts <* spaces
    combined <- P.kwCombine >> 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)]