module LanguageOne where
import qualified PythonAst as Py
import Data.Bifunctor
import Data.Char
import Data.Functor
import qualified Data.Map as Map
import Data.Maybe
import qualified Data.Set as Set
import Text.Parsec
import Text.Parsec.Char
import Text.Parsec.Combinator
import Control.Monad.State

{- Data Types -}
data PossibleType = List | Any deriving Eq

data SelectorMarker = None | Remove

data Op
    = Add
    | Subtract
    | Multiply
    | Divide
    | Insert
    | Concat
    | LessThan
    | LessThanEq
    | GreaterThan
    | GreaterThanEq 
    | Equal
    | NotEqual
    | And
    | Or

data Selector = Selector String Expr

data Expr
    = Var String
    | IntLiteral Int
    | ListLiteral [Expr]
    | Split Expr [Selector] Expr
    | IfElse Expr Expr Expr
    | BinOp Op Expr Expr
    | FunctionCall Expr [Expr]
    | LengthOf Expr
    | Random
    | Access Expr Expr SelectorMarker
    | Parameter Int

data Function = Function String [String] Expr

data Prog = Prog [Function]

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

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

parseVar :: Parser String
parseVar =
    do
        c <- satisfy (\c -> (isLetter c && isLower c) || c == '_')
        cs <- many (satisfy isLetter <|> digit)
        spaces
        let var = c:cs
        if var `elem` ["if", "then", "else", "rand"]
        then fail "reserved"
        else return var

parseKwIf :: Parser ()
parseKwIf = string "if" $> ()

parseKwThen :: Parser ()
parseKwThen = string "then" $> ()

parseKwElse :: Parser ()
parseKwElse = string "else" $> ()

parseKwRand :: Parser Expr
parseKwRand = string "rand" $> Random

parseThis :: Parser Expr
parseThis =
    do
        char '&'
        contextNum <- getState
        spaces
        return (Var $ "context_" ++ show contextNum)

parseList :: Parser Expr
parseList = ListLiteral <$>
    do
        char '[' >> spaces
        es <- sepBy parseExpr (char ',' >> spaces)
        spaces >> char ']' >> spaces
        return es

parseSplit :: Parser Expr
parseSplit =
    do
        char '~' >> spaces
        e <- parseExpr
        spaces >> string "->"
        spaces >> char '{'
        contextNum <- getState
        putState $ return $ 1 + fromMaybe (-1) contextNum
        es <- many1 (spaces >> parseSelector)
        putState contextNum
        spaces >> char '}' >> spaces >> string "->" >> spaces
        e' <- parseExpr
        spaces
        return $ Split e es e'

parseSelectorMarker :: Parser SelectorMarker
parseSelectorMarker = (char '!' >> return Remove) <|> return None

parseSelector :: Parser Selector
parseSelector =
    do
        name <- parseVar
        spaces >> string "<-" >> spaces
        expr <- parseExpr
        spaces
        return $ Selector name expr

parseIfElse :: Parser Expr
parseIfElse =
    do
        parseKwIf >> spaces
        ec <- parseExpr
        spaces >> parseKwThen >> spaces
        et <- parseExpr
        spaces >> parseKwElse >> spaces
        ee <- parseExpr
        spaces
        return $ IfElse ec et ee

parseLength :: Parser Expr
parseLength =
    do
        char '|' >> spaces
        e <- parseExpr
        spaces >> char '|' >> spaces
        return $ LengthOf e

parseParameter :: Parser Expr
parseParameter =
    do
        char '#'
        d <- digit
        spaces
        return $ Parameter $ read [d]

parseParenthesized :: Parser Expr
parseParenthesized =
    do
        char '(' >> spaces
        e <- parseExpr
        spaces >> char ')' >> spaces
        return e

parseBasicExpr :: Parser Expr
parseBasicExpr = choice
    [ IntLiteral <$> parseInt
    , parseThis
    , parseList
    , parseSplit
    , parseLength
    , parseParameter
    , parseParenthesized
    , Var <$> try parseVar
    , parseKwRand
    , parseIfElse
    ] 

parsePostfix :: Parser (Expr -> Expr)
parsePostfix = parsePostfixAccess <|> parsePostfixCall

parsePostfixAccess :: Parser (Expr -> Expr)
parsePostfixAccess =
    do
        char '[' >> spaces
        e <- parseExpr
        spaces >> char ']' >> spaces
        marker <- parseSelectorMarker
        spaces
        return $ \e' -> Access e' e marker

parsePostfixCall :: Parser (Expr -> Expr)
parsePostfixCall =
    do
        char '(' >> spaces
        es <- sepBy parseExpr (char ',' >> spaces)
        char ')' >> spaces
        return $ flip FunctionCall es

parsePostfixedExpr :: Parser Expr
parsePostfixedExpr =
    do
        eb <- parseBasicExpr
        spaces
        ps <- many parsePostfix
        return $ foldl (flip ($)) eb ps

parseOp :: String -> Op -> Parser Op
parseOp s o = try (string s) >> return o

parseLevel :: Parser Expr -> Parser Op -> Parser Expr
parseLevel pe po =
    do
        start <- pe
        spaces
        ops <- many $ try $ do
            op <- po
            spaces
            val <- pe
            spaces
            return (op, val)
        spaces
        return $ foldl (\l (o, r) -> BinOp o l r) start ops

parseExpr :: Parser Expr
parseExpr = foldl parseLevel parsePostfixedExpr
    [ parseOp "*" Multiply, parseOp "/" Divide
    , parseOp "+" Add, parseOp "-" Subtract
    , parseOp "<<" Insert
    , parseOp "++" Concat
    , parseOp "<=" LessThanEq <|> parseOp ">=" GreaterThanEq <|>
        parseOp "<" LessThan <|> parseOp ">" GreaterThan <|>
        parseOp "==" Equal <|> parseOp "!=" NotEqual
    , parseOp "&&" And <|> parseOp "||" Or
    ]

parseFunction :: Parser Function
parseFunction =
    do
        name <- parseVar
        spaces >> char '(' >> spaces
        vs <- sepBy parseVar (char ',' >> spaces)
        spaces >> char ')' >> spaces >> char '=' >> spaces
        body <- parseExpr
        spaces
        return $ Function name vs body

parseProg :: Parser Prog
parseProg = Prog <$> sepBy1 parseFunction (char ';' >> spaces)

parse :: SourceName -> String -> Either ParseError Prog
parse = runParser parseProg Nothing 

{- "Type" checker -}
mergePossibleType :: PossibleType -> PossibleType -> PossibleType
mergePossibleType List _ = List
mergePossibleType _ List = List
mergePossibleType _ _ = Any

getPossibleType :: String -> Expr -> PossibleType
getPossibleType s (Var s') = if s == s' then List else Any
getPossibleType _ (ListLiteral _) = List
getPossibleType s (Split _ _ e) = getPossibleType s e
getPossibleType s (IfElse i t e) =
    foldl1 mergePossibleType $ map (getPossibleType s) [i, t, e]
getPossibleType _ (BinOp Insert _ _) = List
getPossibleType _ (BinOp Concat _ _) = List
getPossibleType _ _ = Any

{- Translator -}
type Translator = Control.Monad.State.State (Map.Map String [String], Int)

currentTemp :: Translator String
currentTemp = do
    t <- gets snd
    return $ "temp" ++ show t

incrementTemp :: Translator String
incrementTemp = do
    modify (second (+1))
    currentTemp

hasLambda :: Expr -> Bool
hasLambda (ListLiteral es) = any hasLambda es
hasLambda (Split e ss r) =
    hasLambda e || any (\(Selector _ e') -> hasLambda e') ss || hasLambda r
hasLambda (IfElse i t e) = hasLambda i || hasLambda t || hasLambda e
hasLambda (BinOp o l r) = hasLambda l || hasLambda r
hasLambda (FunctionCall e es) = any hasLambda $ e : es
hasLambda (LengthOf e) = hasLambda e
hasLambda (Access e _ _) = hasLambda e
hasLambda Parameter{} = True
hasLambda _ = False

translate :: Prog -> [Py.PyStmt]
translate p = fst $ runState (translateProg p) (Map.empty, 0)

translateProg :: Prog -> Translator [Py.PyStmt]
translateProg (Prog fs) = concat <$> traverse translateFunction fs

translateFunction :: Function -> Translator [Py.PyStmt]
translateFunction (Function n ps ex) = do
    let createIf p = Py.BinOp Py.Equal (Py.Var p) (Py.ListLiteral [])
    let createReturn p = Py.IfElse (createIf p) [Py.Return (Py.Var p)] [] Nothing
    let fastReturn = [createReturn p | p <- take 1 ps, getPossibleType p ex == List]
    (ss, e) <- translateExpr ex
    return $ return $ Py.FunctionDef n ps $ fastReturn ++ ss ++ [Py.Return e]

translateSelector :: Selector -> Translator Py.PyStmt
translateSelector (Selector n e) =
    let
        cacheCheck = Py.NotIn (Py.StrLiteral n) (Py.Var "cache")
        cacheAccess = Py.Access (Py.Var "cache") [Py.StrLiteral n]
        cacheSet = Py.Assign (Py.AccessPat (Py.Var "cache") [Py.StrLiteral n])
        body e' = [ Py.IfElse cacheCheck [cacheSet e'] [] Nothing, Py.Return cacheAccess]
    in
        do
            (ss, e') <- translateExpr e
            vs <- gets fst
            let callPrereq p = Py.Standalone $ Py.FunctionCall (Py.Var p) []
            let prereqs = maybe [] (map callPrereq) $ Map.lookup n vs
            return $ Py.FunctionDef n [] $ ss ++ prereqs ++ body e'

translateExpr :: Expr -> Translator ([Py.PyStmt], Py.PyExpr)
translateExpr (Var s) = do
    vs <- gets fst
    let sVar = Py.Var s
    let expr = if Map.member s vs then Py.FunctionCall sVar [] else sVar
    return ([], expr) 
translateExpr (IntLiteral i) = return ([], Py.IntLiteral i)
translateExpr (ListLiteral l) = do
    tl <- mapM translateExpr l
    return (concatMap fst tl, Py.ListLiteral $ map snd tl)
translateExpr (Split e ss e') = do
    vs <- gets fst
    let cacheAssign = Py.Assign (Py.VarPat "cache") (Py.DictLiteral [])
    let cacheStmt = [ cacheAssign | Map.size vs == 0 ]
    let vnames = map (\(Selector n es) -> n) ss
    let prereqs = snd $ foldl (\(ds, m) (Selector n es) -> (n:ds, Map.insert n ds m)) ([], Map.empty) ss
    modify $ first $ Map.union prereqs
    fs <- mapM translateSelector ss
    (sts, te) <- translateExpr e'
    modify $ first $ const vs
    return (cacheStmt ++ fs ++ sts, te)
translateExpr (IfElse i t e) = do
    temp <- incrementTemp
    let tempPat = Py.VarPat temp
    (ists, ie) <- translateExpr i
    (tsts, te) <- translateExpr t
    (ests, ee) <- translateExpr e
    let thenSts = tsts ++ [Py.Assign tempPat te]
    let elseSts = ests ++ [Py.Assign tempPat ee]
    let newIf = Py.IfElse ie thenSts [] $ Just elseSts
    return (ists ++ [newIf], Py.Var temp)
translateExpr (BinOp o l r) = do
    (lsts, le) <- translateExpr l
    (rsts, re) <- translateExpr r
    (opsts, oe) <- translateOp o le re
    return (lsts ++ rsts ++ opsts, oe)
translateExpr (FunctionCall f ps) = do
    (fsts, fe) <- translateExpr f
    tps <- mapM translateExpr ps
    return (fsts ++ concatMap fst tps, Py.FunctionCall fe $ map snd tps)
translateExpr (LengthOf e) =
    second (Py.FunctionCall (Py.Var "len") . return) <$> translateExpr e
translateExpr (Access e Random m) = do
    temp <- incrementTemp
    (sts, ce) <- translateExpr e
    let lenExpr = Py.FunctionCall (Py.Var "len") [Py.Var temp]
    let randExpr = Py.FunctionCall (Py.Var "randint") [ Py.IntLiteral 0,  lenExpr ]
    return (sts, singleAccess ce randExpr m)
translateExpr (Access c i m) = do
    (csts, ce) <- translateExpr c
    (ists, ie) <- translateExpr i
    temp <- incrementTemp
    if hasLambda i
    then return (csts ++ ists ++ [createFilterLambda temp ie m], Py.FunctionCall (Py.Var temp) [ce])
    else return (csts ++ ists, singleAccess ce ie m)
translateExpr (Parameter i) = return $ ([], Py.Var $ "arg" ++ show i)
translateExpr _ = fail "Invalid expression"

singleAccess :: Py.PyExpr -> Py.PyExpr -> SelectorMarker -> Py.PyExpr
singleAccess c i None = Py.Access c [i]
singleAccess c i Remove = Py.FunctionCall (Py.Member c "pop") [i]

createFilterLambda :: String -> Py.PyExpr -> SelectorMarker -> Py.PyStmt
createFilterLambda s e None = Py.FunctionDef s ["arg"]
    [ Py.Assign (Py.VarPat "out") (Py.ListLiteral [])
    , Py.For (Py.VarPat "arg0") (Py.Var "arg")
        [ Py.IfElse e
            [ Py.Standalone $ Py.FunctionCall (Py.Member (Py.Var "out") "append")
                [ Py.Var "arg0" ]
            ]
            []
            Nothing
        ]
    , Py.Return $ Py.Var "out"
    ]
createFilterLambda s e Remove = Py.FunctionDef s ["arg"]
    [ Py.Assign (Py.VarPat "i") $ Py.IntLiteral 0
    , Py.Assign (Py.VarPat "out") (Py.ListLiteral [])
    , Py.While (Py.BinOp Py.LessThan (Py.Var "i") $ Py.FunctionCall (Py.Var "len") [Py.Var "arg"])
        [ Py.IfElse e
            [ Py.Standalone $ Py.FunctionCall (Py.Member (Py.Var "out") "append")
                [ singleAccess (Py.Var "arg") (Py.Var "i") Remove
                ]
            ]
            []
            Nothing
        , Py.Assign (Py.VarPat "i") (Py.BinOp Py.Add (Py.Var "i") (Py.IntLiteral 1))
        ]
    , Py.Return $ Py.Var "out"
    ]

translateOp :: Op -> Py.PyExpr -> Py.PyExpr -> Translator ([Py.PyStmt], Py.PyExpr)
translateOp Add l r = return ([], Py.BinOp Py.Add l r)
translateOp Subtract l r = return ([], Py.BinOp Py.Subtract l r)
translateOp Multiply l r = return ([], Py.BinOp Py.Multiply l r)
translateOp Divide l r = return ([], Py.BinOp Py.Divide l r)
translateOp LessThan l r = return ([], Py.BinOp Py.LessThan l r)
translateOp LessThanEq l r = return ([], Py.BinOp Py.LessThanEq l r)
translateOp GreaterThan l r = return ([], Py.BinOp Py.GreaterThan l r)
translateOp GreaterThanEq  l r = return ([], Py.BinOp Py.GreaterThanEq  l r)
translateOp Equal l r = return ([], Py.BinOp Py.Equal l r)
translateOp NotEqual l r = return ([], Py.BinOp Py.NotEqual l r)
translateOp And l r = return ([], Py.BinOp Py.And l r)
translateOp Or l r = return ([], Py.BinOp Py.Or l r)
translateOp Concat l r = return ([], Py.BinOp Py.Add l r)
translateOp Insert l r = do
    temp <- incrementTemp
    let assignStmt = Py.Assign (Py.VarPat temp) l
    let appendFunc = Py.Member (Py.Var temp) "append"
    let insertStmt = Py.Standalone $ Py.FunctionCall appendFunc [r]
    return ([assignStmt, insertStmt], Py.Var temp)