{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
module Show where

import Control.Monad.State.Strict

import AST
import ASTfunc


showExpr :: Expr env t -> String
showExpr e = showsExpr e ""

showsExpr :: Expr env t -> ShowS
showsExpr e = evalState (showsExpr' 0 [] e) 1

showsExpr' :: Int -> [String] -> Expr env t -> State Int ShowS
showsExpr' d env (Expr e) =
    showsPExpr' showsExpr' (\case Expr (Pair a b) -> Just (a, b) ; _ -> Nothing) etypeOf d env e

showsPExpr' :: (forall env' t'. Int -> [String] -> expr env' t' -> State Int ShowS)
            -> (forall env' t1 t2. expr env' (t1, t2) -> Maybe (expr env' t1, expr env' t2))
            -> (forall env' t'. expr env' t' -> Tup STy t')
            -> Int -> [String] -> PExpr expr env t -> State Int ShowS
showsPExpr' _ _ _ _ _ (Const t x) = return (showsScalar t x)
showsPExpr' f _ _ _ env (Pair e1 e2) = do
    s1 <- f 0 env e1
    s2 <- f 0 env e2
    return (showParen True (s1 . showString ", " . s2))
showsPExpr' _ _ _ _ _ Nil = return (showString "()")
showsPExpr' f pf tf d env (Prim op e) = case (showOper op, tf e) of
    (Infix opd s, T2 _ _) | Just (e1, e2) <- pf e -> do
        s1 <- f (opd + 1) env e1
        s2 <- f (opd + 1) env e2
        return (showParen (d > opd) (s1 . showString (" " ++ s ++ " ") . s2))
    (Infix _ s, _) -> do
        s1 <- f 11 env e
        return (showParen (d > 10) (showString (s ++ " ") . s1))
    (PrefixFun s, _) -> do
        s1 <- f 11 env e
        return (showParen (d > 10) (showString (s ++ " ") . s1))
showsPExpr' _ _ _ _ env (Var (V _ i)) =
    case drop (idxToInt i) env of
      s:_ -> return (showString s)
      [] -> return (showString ("tUP" ++ show (idxToInt i - length env + 1)))
showsPExpr' f _ _ d env (Let lhs rhs e) = do
    (lhss, envf) <- nameifyLHS lhs
    s1 <- f 0 env rhs
    s2 <- f 0 (envf env) e
    return (showParen (d > 0) (showString ("let " ++ lhss ++ " = ") . s1 . showString " in " . s2))

nameifyLHS :: LHS s t env env' -> State Int (String, [String] -> [String])
nameifyLHS (L0 _) = return ("_", id)
nameifyLHS (L1 _) = do
    seed <- get
    put (seed + 1)
    let name = 't' : show seed
    return (name, (name :))
nameifyLHS (L2 l1 l2) = do
    (s1, f1) <- nameifyLHS l1
    (s2, f2) <- nameifyLHS l2
    return ("(" ++ s1 ++ "," ++ s2 ++ ")", f2 . f1)

data ShownOper t = Infix Int String | PrefixFun String

showOper :: Oper (t -> t') -> ShownOper t
showOper (Add _) = Infix 4 "+"
showOper (Mul _) = Infix 5 "*"
showOper Round = PrefixFun "round"

showsScalar :: STy t -> t -> ShowS
showsScalar TInt = shows
showsScalar TFloat = shows