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

parseVar :: Parser String
parseVar = P.var ["if", "then", "else", "var"]

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
        P.kwIf >> spaces
        ec <- parseExpr
        spaces >> P.kwThen >> spaces
        et <- parseExpr
        spaces >> P.kwElse >> 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 <$> P.int
    , parseThis
    , parseList
    , parseSplit
    , parseLength
    , parseParameter
    , parseParenthesized
    , Var <$> try parseVar
    , P.kwRand $> Random
    , 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

parseExpr :: Parser Expr
parseExpr = P.precedence BinOp parsePostfixedExpr
    [ P.op "*" Multiply, P.op "/" Divide
    , P.op "+" Add, P.op "-" Subtract
    , P.op "<<" Insert
    , P.op "++" Concat
    , try (P.op "<=" LessThanEq) <|> try (P.op ">=" GreaterThanEq) <|>
        P.op "<" LessThan <|> P.op ">" GreaterThan <|>
        P.op "==" Equal <|> P.op "!=" NotEqual
    , P.op "&&" And <|> P.op "||" 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)