{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module ReduceC where

import Control.Monad.Reduce

import Control.Applicative
import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import Data.Functor
import Data.Maybe
import qualified Language.C as C

type Lab = C.Ident

reduceC :: (MonadReduce Lab m) => C.CTranslUnit -> m C.CTranslUnit
reduceC (C.CTranslUnit es ni) = do
  es' <- collect mrCExternalDeclaration es
  pure $ C.CTranslUnit es' ni

mrCExternalDeclaration :: (MonadReduce Lab m) => C.CExternalDeclaration C.NodeInfo -> MaybeT m (C.CExternalDeclaration C.NodeInfo)
mrCExternalDeclaration = \case
  C.CFDefExt fun -> do
    givenWith (funName fun)
    C.CFDefExt <$> rCFunctionDef fun
  C.CDeclExt decl ->
    C.CDeclExt <$> mrCDeclaration decl
  a -> error (show a)
 where
  funName (C.CFunDef _ (C.CDeclr x _ _ _ _) _ _ _) =
    x

mrCDeclarationItem :: (MonadReduce Lab m) => C.CDeclarationItem C.NodeInfo -> MaybeT m (C.CDeclarationItem C.NodeInfo)
mrCDeclarationItem = \case
  C.CDeclarationItem d@(C.CDeclr x _ _ _ _) i e -> do
    givenWith x
    i' <- mtry $ munder i mrCInitializer
    e' <- mtry $ munder e mrCExpression
    pure (C.CDeclarationItem d i' e')
  a -> error (show a)

mrCInitializer :: (MonadReduce Lab m) => C.CInitializer C.NodeInfo -> MaybeT m (C.CInitializer C.NodeInfo)
mrCInitializer = \case
  C.CInitExpr e ni -> mrCExpression e <&> \e' -> C.CInitExpr e' ni
  C.CInitList (C.CInitializerList items) ni -> do
    collectNonEmpty' rmCInitializerListItem items <&> \items' ->
      C.CInitList (C.CInitializerList items') ni
 where
  rmCInitializerListItem (pds, is) = do
    pds' <- lift $ collect rmCPartDesignator pds
    is' <- mrCInitializer is
    pure (pds', is')

  rmCPartDesignator :: (MonadReduce Lab m) => C.CPartDesignator C.NodeInfo -> m (C.CPartDesignator C.NodeInfo)
  rmCPartDesignator = \case
    a -> error (show a)

mrCDeclaration :: (MonadReduce Lab m) => C.CDeclaration C.NodeInfo -> MaybeT m (C.CDeclaration C.NodeInfo)
mrCDeclaration = \case
  C.CDecl spc decl ni -> do
    decl' <- lift $ collect mrCDeclarationItem decl
    case decl' of
      [] -> empty
      decl'' -> pure $ C.CDecl spc decl'' ni
  a -> error (show a)

rCFunctionDef :: (MonadReduce Lab m) => C.CFunctionDef C.NodeInfo -> m (C.CFunctionDef C.NodeInfo)
rCFunctionDef (C.CFunDef spc dec cdecls smt ni) = do
  smt' <- rCStatement smt
  pure $ C.CFunDef spc dec cdecls smt' ni

rCStatement :: (MonadReduce Lab m) => C.CStatement C.NodeInfo -> m (C.CStatement C.NodeInfo)
rCStatement = \case
  C.CCompound is cbi ni -> do
    cbi' <- collect mrCCompoundBlockItem cbi
    pure $ C.CCompound is cbi' ni
  C.CExpr e ni -> do
    e' <- runMaybeT $ munder e mrCExpression
    pure $ C.CExpr e' ni
  C.CIf e s els ni -> do
    e' <- runMaybeT $ mrCExpression e
    s' <- rCStatement s
    els' <- case els of
      Just els' -> do
        pure Nothing <| Just <$> rCStatement els'
      Nothing -> pure Nothing
    case (e', els') of
      (Nothing, Nothing) -> pure s'
      (Just e'', Nothing) -> pure $ C.CIf e'' s' Nothing ni
      (Nothing, Just x) -> pure $ C.CIf zeroExp s' (Just x) ni
      (Just e'', Just x) -> pure $ C.CIf e'' s' (Just x) ni
  C.CFor e1 e2 e3 s ni -> do
    rCStatement s <| do
      e1' <- rCForInit e1
      e2' <- runMaybeT $ munder e2 mrCExpression
      e3' <- runMaybeT $ munder e3 mrCExpression
      s' <- rCStatement s
      pure $ C.CFor e1' e2' e3' s' ni
  C.CReturn e ni -> do
    e' <- case e of
      Nothing -> pure Nothing
      Just e' -> Just <$> zrCExpression e'
    pure $ C.CReturn e' ni
  C.CBreak ni -> pure (C.CBreak ni)
  C.CCont ni -> pure (C.CCont ni)
  C.CLabel i s [] ni ->
    -- todo fix attrs
    splitOn i (rCStatement s) do
      s' <- rCStatement s
      pure $ C.CLabel i s' [] ni
  C.CGoto i ni ->
    -- todo fix attrs
    splitOn i (pure $ C.CExpr Nothing ni) do
      pure $ C.CGoto i ni
  a -> error (show a)
 where
  rCForInit = \case
    C.CForDecl decl -> do
      m <- runMaybeT $ mrCDeclaration decl
      pure $ case m of
        Nothing -> C.CForInitializing Nothing
        Just d' -> C.CForDecl d'
    C.CForInitializing n -> do
      C.CForInitializing <$> runMaybeT (munder n mrCExpression)

orZero :: Maybe (C.CExpression C.NodeInfo) -> C.CExpression C.NodeInfo
orZero = fromMaybe zeroExp

zeroExp :: C.CExpression C.NodeInfo
zeroExp = C.CConst (C.CIntConst (C.cInteger 0) C.undefNode)

zrCExpression :: (MonadReduce Lab m) => C.CExpression C.NodeInfo -> m (C.CExpression C.NodeInfo)
zrCExpression e = orZero <$> runMaybeT (mrCExpression e)

mrCExpression :: (MonadReduce Lab m) => C.CExpression C.NodeInfo -> MaybeT m (C.CExpression C.NodeInfo)
mrCExpression = \case
  C.CVar i ni -> do
    givenThat i
    pure $ C.CVar i ni
  C.CCall e es ni -> do
    e' <- mrCExpression e
    es' <- lift $ traverse zrCExpression es
    pure $ C.CCall e' es' ni
  C.CCond ec et ef ni -> do
    ec' <- mrCExpression ec
    ef' <- mrCExpression ef
    et' <- mtry $ munder et mrCExpression
    pure $ C.CCond ec' et' ef' ni
  C.CBinary o elhs erhs ni -> onBothExpr elhs erhs \lhs rhs ->
    pure $ C.CBinary o lhs rhs ni
  C.CUnary o elhs ni -> do
    lhs <- mrCExpression elhs
    pure $ C.CUnary o lhs ni
  C.CConst c -> do
    -- TODO fix
    pure $ C.CConst c
  C.CCast cd e ni -> do
    -- TODO fix
    cd' <- mrCDeclaration cd
    e' <- mrCExpression e
    pure $ C.CCast cd' e' ni
  C.CAssign op e1 e2 ni -> onBothExpr e1 e2 \e1' e2' ->
    pure $ C.CAssign op e1' e2' ni
  C.CIndex e1 e2 ni -> onBothExpr e1 e2 \e1' e2' ->
    pure $ C.CIndex e1' e2' ni
  C.CMember e i b ni -> do
    givenThat i
    e' <- mrCExpression e
    pure $ C.CMember e' i b ni
  C.CComma items ni -> do
    C.CComma <$> collectNonEmpty' mrCExpression items <*> pure ni
  e -> error (show e)
 where
  onBothExpr elhs erhs = onBoth (mrCExpression elhs) (mrCExpression erhs)

mrCCompoundBlockItem
  :: (MonadReduce Lab m)
  => C.CCompoundBlockItem C.NodeInfo
  -> MaybeT m (C.CCompoundBlockItem C.NodeInfo)
mrCCompoundBlockItem = \case
  C.CBlockStmt s -> empty <| lift (C.CBlockStmt <$> rCStatement s)
  C.CBlockDecl d -> C.CBlockDecl <$> mrCDeclaration d
  a -> error (show a)

mtry :: (Functor m) => MaybeT m a -> MaybeT m (Maybe a)
mtry (MaybeT mt) = MaybeT (Just <$> mt)

mlift :: (Applicative m) => Maybe a -> MaybeT m a
mlift a = MaybeT (pure a)

munder :: (Monad m) => Maybe a -> (a -> MaybeT m b) -> MaybeT m b
munder a mf = mlift a >>= mf