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

module ReduceC where

import Control.Monad.Reduce

import qualified Data.Valuation as Val

import Control.Applicative
import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import Data.Foldable
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
    C.CFDefExt <$> mrCFunctionDef fun
  C.CDeclExt decl ->
    C.CDeclExt <$> mrCDeclaration decl
  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
    mapM_ cCDeclarationSpecifier spc
    decl' <- lift $ collect mrCDeclarationItem decl
    case decl' of
      [] -> empty
      decl'' -> pure $ C.CDecl spc decl'' ni
  a -> error (show a)

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

cCDeclarationItem :: (MonadReduce Lab m) => C.CDeclarationItem C.NodeInfo -> MaybeT m ()
cCDeclarationItem = \case
  C.CDeclarationItem d i e -> do
    munder i cCInitializer
    munder e cCExpression
    cCDeclr d
  a -> error (show a)

cCDeclaration :: (MonadReduce Lab m) => C.CDeclaration C.NodeInfo -> MaybeT m ()
cCDeclaration = \case
  C.CDecl spc decl _ -> do
    forM_ spc cCDeclarationSpecifier
    mapM_ cCDeclarationItem decl
  a -> error (show a)

cCExpression :: (MonadReduce Lab m) => C.CExpression C.NodeInfo -> MaybeT m ()
cCExpression e =
  -- TODO not optimal, create version that only checks for identifiers  .
  void $ mrCExpression e

cCDeclarationSpecifier :: (MonadReduce Lab m) => C.CDeclarationSpecifier C.NodeInfo -> MaybeT m ()
cCDeclarationSpecifier = \case
  C.CTypeSpec t -> cCTypeSpecifier t
  C.CStorageSpec _ -> pure ()
  C.CTypeQual t -> cCTypeQualifier t
  C.CFunSpec _ -> pure ()
  C.CAlignSpec (C.CAlignAsType t _) -> cCDeclaration t
  C.CAlignSpec (C.CAlignAsExpr t _) -> cCExpression t

cCTypeQualifier :: (MonadReduce Lab m) => C.CTypeQualifier C.NodeInfo -> MaybeT m ()
cCTypeQualifier = \case
  C.CAttrQual a -> cCAttr a
  _ -> pure ()

cCTypeSpecifier :: (MonadReduce Lab m) => C.CTypeSpecifier C.NodeInfo -> MaybeT m ()
cCTypeSpecifier = \case
  C.CVoidType _ -> pure ()
  C.CCharType _ -> pure ()
  C.CShortType _ -> pure ()
  C.CIntType _ -> pure ()
  C.CLongType _ -> pure ()
  C.CFloatType _ -> pure ()
  C.CDoubleType _ -> pure ()
  C.CSignedType _ -> pure ()
  C.CUnsigType _ -> pure ()
  C.CBoolType _ -> pure ()
  C.CComplexType _ -> pure ()
  C.CInt128Type _ -> pure ()
  -- C.CUInt128Type a -> pure ()
  C.CFloatNType{} -> pure ()
  C.CTypeDef i _ -> do
    givenThat (Val.is i)
    pure ()
  (C.CTypeOfExpr e _) -> cCExpression e
  (C.CTypeOfType t _) -> cCDeclaration t
  (C.CAtomicType t _) -> cCDeclaration t
  a@(C.CSUType _ _) -> error (show a)
  a@(C.CEnumType _ _) -> error (show a)

cCInitializer :: (MonadReduce Lab m) => C.CInitializer C.NodeInfo -> MaybeT m ()
cCInitializer = void . mrCInitializer

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)

mrCFunctionDef :: (MonadReduce Lab m) => C.CFunctionDef C.NodeInfo -> MaybeT m (C.CFunctionDef C.NodeInfo)
mrCFunctionDef (C.CFunDef spc dec cdecls smt ni) = do
  smt' <- lift $ rCStatement smt
  mapM_ cCDeclaration cdecls
  mapM_ cCDeclarationSpecifier spc
  cCDeclr dec
  pure $ C.CFunDef spc dec cdecls smt' ni

cCDeclr :: (MonadReduce Lab m) => C.CDeclarator C.NodeInfo -> MaybeT m ()
cCDeclr (C.CDeclr x dd _ _ _) = do
  mapM_ cCDerivedDeclarator dd
  givenWith (Val.is <$> x)
 where
  cCDerivedDeclarator = \case
    C.CPtrDeclr ts _ -> mapM_ cCTypeQualifier ts
    C.CArrDeclr ts as _ -> do
      mapM_ cCTypeQualifier ts
      case as of
        C.CNoArrSize _ -> pure ()
        C.CArrSize _ e -> cCExpression e
    C.CFunDeclr f attr _ -> do
      mapM_ cCAttr attr
      cCFunParams f
  cCFunParams = \case
    C.CFunParamsOld o -> mapM_ (givenThat . Val.is) o
    C.CFunParamsNew o _ -> mapM_ cCDeclaration o

cCAttr :: (MonadReduce Lab m) => C.CAttribute C.NodeInfo -> MaybeT m ()
cCAttr (C.CAttr i e _) = do
  mapM_ cCExpression e
  givenThat (Val.is i)

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 (Val.is i) (rCStatement s) do
      s' <- rCStatement s
      pure $ C.CLabel i s' [] ni
  C.CGoto i ni ->
    -- todo fix attrs
    splitOn (Val.is i) (pure $ C.CExpr Nothing ni) do
      pure $ C.CGoto i ni
  C.CWhile e s dow ni -> do
    e' <- zrCExpression e
    s' <- rCStatement s
    pure $ C.CWhile e' s' dow 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 (Val.is 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 (Val.is 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