module PythonGen where
import PythonAst
import Data.List
import Data.Bifunctor
import Data.Maybe

indent :: String -> String
indent = ("    " ++)

stmtBlock :: [PyStmt] -> [String]
stmtBlock = concatMap translateStmt

block :: String -> [String] -> [String]
block s ss = (s ++ ":") : map indent ss

prefix :: String -> PyExpr -> [PyStmt] -> [String]
prefix s e sts = block (s ++ " " ++ translateExpr e) $ stmtBlock sts

if_ :: PyExpr -> [PyStmt] -> [String]
if_ = prefix "if"

elif :: PyExpr -> [PyStmt] -> [String]
elif = prefix "elif"

else_ :: [PyStmt] -> [String]
else_ = block "else" . stmtBlock

while :: PyExpr -> [PyStmt] -> [String]
while = prefix "while"

parenth :: String -> String
parenth s = "(" ++ s ++ ")"

translateStmt :: PyStmt -> [String]
translateStmt (Assign p e) = [translatePat p ++ " = " ++ translateExpr e]
translateStmt (IfElse i t es e) =
    if_ i t ++ concatMap (uncurry elif) es ++ maybe [] else_ e
translateStmt (While c t) = while c t
translateStmt (For x in_ b) = block head body
    where
        head = "for " ++ translatePat x ++ " in " ++ translateExpr in_
        body = stmtBlock b
translateStmt (FunctionDef s ps b) = block head body
    where
        head = "def " ++ s ++ "(" ++ intercalate "," ps ++ ")"
        body = stmtBlock b
translateStmt (Return e) = ["return " ++ translateExpr e]
translateStmt (Standalone e) = [translateExpr e]
translateStmt (Import s) = ["import " ++ s]
translateStmt (FromImport s ss) =
    ["from " ++ s ++ " import " ++ intercalate "," ss]
translateStmt (Nonlocal vs) = 
    ["nonlocal " ++ intercalate "," vs]

precedence :: PyBinOp -> Int
precedence Add = 3
precedence Subtract = 3
precedence Multiply = 4
precedence Divide = 4
precedence FloorDiv = 4
precedence LessThan = 2
precedence LessThanEq = 2
precedence GreaterThan = 2
precedence GreaterThanEq = 2
precedence Equal = 2
precedence NotEqual = 2
precedence And = 1
precedence Or = 0

opString :: PyBinOp -> String
opString Add = "+"
opString Subtract = "-"
opString Multiply = "*"
opString Divide = "/"
opString FloorDiv = "//"
opString LessThan = "<"
opString LessThanEq = "<="
opString GreaterThan = ">"
opString GreaterThanEq = ">="
opString Equal = "=="
opString NotEqual = "!="
opString And = " and "
opString Or = " or "

translateOp :: PyBinOp -> PyBinOp -> PyExpr -> String
translateOp o o' =
    if precedence o > precedence o'
    then parenth . translateExpr
    else translateExpr

dictMapping :: PyExpr -> PyExpr -> String
dictMapping f t = translateExpr f ++ ": " ++ translateExpr t

list :: String -> String -> [PyExpr] -> String
list o c es = o ++ intercalate ", " (map translateExpr es) ++ c

translateExpr :: PyExpr -> String
translateExpr (BinOp o l@(BinOp o1 _ _) r@(BinOp o2 _ _)) =
    translateOp o o1 l ++ opString o ++ translateOp o o2 r
translateExpr (BinOp o l@(BinOp o1 _ _) r) =
    translateOp o o1 l ++ opString o ++ translateExpr r
translateExpr (BinOp o l r@(BinOp o2 _ _)) =
    translateExpr l ++ opString o ++ translateOp o o2 r
translateExpr (BinOp o l r) =
    translateExpr l ++ opString o ++ translateExpr r
translateExpr (IntLiteral i) = show i
translateExpr (StrLiteral s) = "\"" ++ s ++ "\""
translateExpr (BoolLiteral b) = if b then "true" else "false"
translateExpr (ListLiteral l) = list "[" "]" l
translateExpr (DictLiteral l) =
    "{" ++ intercalate ", " (map (uncurry dictMapping) l) ++ "}"
translateExpr (Lambda ps e) = parenth (head ++ ": " ++ body)
    where
        head = "lambda " ++ intercalate ", " (map translatePat ps)
        body = translateExpr e
translateExpr (Var s) = s
translateExpr (TupleLiteral es) = list "(" ")" es
translateExpr (FunctionCall f ps) = translateExpr f ++ list "(" ")" ps
translateExpr (Access (Var s) e) = s ++ list "[" "]" e
translateExpr (Access e@Access{} i) = translateExpr e ++ list "[" "]" i
translateExpr (Access e i) = "(" ++ translateExpr e ++ ")" ++ list "[" "]" i
translateExpr (Ternary c t e) =
    translateExpr t ++ " if " ++ translateExpr c ++ " else " ++ translateExpr e
translateExpr (Member (Var s) m) = s ++ "." ++ m
translateExpr (Member e@Member{} m) = translateExpr e ++ "." ++ m
translateExpr (Member e m) = "(" ++ translateExpr e ++ ")." ++ m
translateExpr (In m c) =
    "(" ++ translateExpr m ++ ") in (" ++ translateExpr c ++ ")"
translateExpr (NotIn m c) =
    "(" ++ translateExpr m ++ ") not in (" ++ translateExpr c ++ ")"
translateExpr (Slice l r) =
    maybe [] (parenth . translateExpr) l ++ ":" ++ maybe [] (parenth . translateExpr) r

translatePat :: PyPat -> String
translatePat (VarPat s) = s
translatePat IgnorePat = "_"
translatePat (TuplePat ps) =
    "(" ++ intercalate "," (map translatePat ps) ++ ")"
translatePat (AccessPat e es) = translateExpr (Access e es)

translate :: [PyStmt] -> String
translate = intercalate "\n" . concatMap translateStmt