{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module ReduceC where

import Control.Monad.Reader
import Control.Monad.Reduce
import Data.Data
import Data.Foldable
import qualified Data.Map.Strict as Map
import qualified Language.C as C

data Context = Context
  { keepMain :: !Bool
  , typeDefs :: !(Map.Map C.Ident [C.CDeclarationSpecifier C.NodeInfo])
  }

defaultReduceC :: (CReducible a, MonadReduce String m) => a -> m a
defaultReduceC a = runReaderT (reduceC a) defaultContext

addTypeDefs :: [C.Ident] -> [C.CDeclarationSpecifier C.NodeInfo] -> Context -> Context
addTypeDefs ids cs Context{..} =
  Context
    { typeDefs =
        foldl' (\a i -> Map.insert i cs a) typeDefs ids
    , ..
    }

defaultContext :: Context
defaultContext =
  Context
    { keepMain = True
    , typeDefs = Map.empty
    }

class CReducible a where
  reduceC :: (MonadReduce String m) => a -> ReaderT Context m a

instance CReducible C.CTranslUnit where
  reduceC (C.CTranslUnit es ni) = do
    es' <- reduceDeclarations es
    pure $ C.CTranslUnit es' ni
   where
    reduceDeclarations = \case
      [] -> pure []
      r : rest -> reduceCExternalDeclaration r (reduceDeclarations rest)

    reduceCExternalDeclaration r cont = do
      shouldKeepMain <- asks keepMain
      case r of
        C.CFDefExt fun
          | shouldKeepMain && maybe False (("main" ==) . C.identToString) (functionName fun) -> do
              r' <- C.CFDefExt <$> reduceC fun
              (r' :) <$> cont
          | otherwise ->
              split ("remove function " <> show (functionName fun)) cont do
                r' <- C.CFDefExt <$> reduceC fun
                (r' :) <$> cont
        C.CDeclExt result ->
          case result of
            -- A typedef
            C.CDecl (C.CStorageSpec (C.CTypedef _) : rst) decl _ -> do
              let ids = identifiers decl
              split
                ("inline typedefs " <> show ids)
                (local (addTypeDefs ids rst) cont)
                ((r :) <$> cont)
            a -> error (show a)
        _r -> error (show r)

instance CReducible C.CFunDef where
  reduceC r = do
    C.CFunDef spc dec cdecls smt ni <- inlineTypeDefs r
    pure $ C.CFunDef spc dec cdecls smt ni

inlineTypeDefs :: forall d m. (Data d, MonadReader Context m) => d -> m d
inlineTypeDefs r = do
  case eqT @d @[C.CDeclarationSpecifier C.NodeInfo] of
    Just Refl -> do
      res' :: [[C.CDeclarationSpecifier C.NodeInfo]] <- forM r \case
        a@(C.CTypeSpec (C.CTypeDef idx _)) -> do
          res <- asks (Map.lookup idx . typeDefs)
          case res of
            Just args -> pure args
            Nothing -> pure [a]
        a -> pure [a]
      pure (fold res')
    Nothing ->
      gmapM inlineTypeDefs r

-- instance CReducible C.CExtDecl where
--  reduceC (C.CFunDef spc dec cdecls smt ni) = do
--    pure $ C.CFunDef spc dec cdecls smt ni

identifiers :: forall a. (Data a) => a -> [C.Ident]
identifiers d = case cast d of
  Just l -> [l]
  Nothing -> concat $ gmapQ identifiers d

functionName :: C.CFunctionDef C.NodeInfo -> Maybe C.Ident
functionName = \case
  C.CFunDef _ (C.CDeclr ix _ _ _ _) _ _ _ -> ix

isMain :: C.CFunctionDef C.NodeInfo -> Bool
isMain (C.CFunDef _ (C.CDeclr (Just i) _ _ _ _) _ _ _) =
  C.identToString i == "main"
isMain _ow = False

-- instance CReducible C.CDeclaration where
--   reduce = \case
--     C.CDecl spc@(C.CStorageSpec (C.CTypedef _) : rst) decl ni -> do
--       decl' <-
--         decl & collectNonEmpty' \case
--           C.CDeclarationItem d Nothing Nothing -> do
--             let (x, _) = cDeclaratorIdentifiers d
--             case x of
--               Just x' ->
--                 splitOn
--                   (Val.is x')
--                   ( do
--                       modify (Map.insert x' (Type rst))
--                       mzero
--                   )
--                   (pure $ C.CDeclarationItem d Nothing Nothing)
--               Nothing ->
--                 pure $ C.CDeclarationItem d Nothing Nothing
--           a -> error (show a)
--       pure (C.CDecl spc decl' ni)
--     C.CDecl spc@[C.CTypeSpec (C.CTypeDef i ni')] decl ni -> do
--       x <- gets (Map.lookup i)
--       case x of
--         Just (Type rst) -> do
--           decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers rst) decl
--           pure $ C.CDecl rst decl' ni
--         Nothing -> do
--           decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers spc) decl
--           pure $ C.CDecl spc decl' ni
--     C.CDecl spc decl ni -> do
--       decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers spc) decl
--       pure $ C.CDecl spc decl' ni
--     a -> error (show a)
--    where
--     reduceCDeclarationItem rq' = \case
--       C.CDeclarationItem d i e -> do
--         let (fn, reqs) = cDeclaratorIdentifiers d
--         case fn of
--           Just fn' ->
--             conditionalGivenThat (rq' <> reqs) (Val.is fn')
--           Nothing ->
--             mapM_ (givenThat . Val.is) (rq' <> reqs)
--
--         i' <- optional do
--           liftMaybe i >>= reduce @C.CInitializer
--         e' <- optional do
--           liftMaybe e >>= reduce @C.CExpression
--
--         pure (C.CDeclarationItem d i' e')
--       a -> error (show a)

-- import Control.Monad.Reduce
--
-- import qualified Data.Valuation as Val
--
-- import Control.Applicative
-- import Control.Monad.State
-- import Control.Monad.Trans.Maybe
-- import Data.Data
-- import Data.Function
-- import Data.Functor
-- import qualified Data.Map.Strict as Map
-- import Data.Maybe (catMaybes)
-- import qualified Language.C as C

-- type Lab = C.Ident
--
-- data LabInfo
--   = Type [C.CDeclarationSpecifier C.NodeInfo]
--
-- type CState = Map.Map Lab LabInfo
--
-- reduceC :: (MonadReduce Lab m, MonadState CState m) => C.CTranslUnit -> m C.CTranslUnit
-- reduceC (C.CTranslUnit es ni) = do
--   es' <- collect reduceCExternalDeclaration es
--   pure $ C.CTranslUnit es' ni
--  where
--   reduceCExternalDeclaration = \case
--     C.CFDefExt fun -> do
--       C.CFDefExt <$> reduce @C.CFunctionDef fun
--     C.CDeclExt decl ->
--       C.CDeclExt <$> reduce @C.CDeclaration decl
--     a -> error (show a)
--
-- identifiers :: forall a. (Data a) => a -> [Lab]
-- identifiers d = case cast d of
--   Just l -> [l]
--   Nothing -> concat $ gmapQ identifiers d
--
-- type Reducer m a = a -> m a
--
-- class CReducible c where
--   reduce :: (MonadReducePlus Lab m, MonadState CState m) => Reducer m (c C.NodeInfo)
--
-- cDeclaratorIdentifiers :: C.CDeclarator C.NodeInfo -> (Maybe Lab, [Lab])
-- cDeclaratorIdentifiers (C.CDeclr mi dd _ la _) =
--   (mi, identifiers dd <> identifiers la)
--
-- instance CReducible C.CFunctionDef where
--   reduce (C.CFunDef spc dec cdecls smt ni) = do
--     let (fn, ids) = cDeclaratorIdentifiers dec
--     let requirements = identifiers spc <> identifiers cdecls <> ids
--     case fn of
--       Just fn' ->
--         conditionalGivenThat requirements (Val.is fn')
--       Nothing ->
--         mapM_ (givenThat . Val.is) requirements
--     smt' <- reduce @C.CStatement smt
--     pure $ C.CFunDef spc dec cdecls smt' ni
--
-- instance CReducible C.CDeclaration where
--   reduce = \case
--     C.CDecl spc@(C.CStorageSpec (C.CTypedef _) : rst) decl ni -> do
--       decl' <-
--         decl & collectNonEmpty' \case
--           C.CDeclarationItem d Nothing Nothing -> do
--             let (x, _) = cDeclaratorIdentifiers d
--             case x of
--               Just x' ->
--                 splitOn
--                   (Val.is x')
--                   ( do
--                       modify (Map.insert x' (Type rst))
--                       mzero
--                   )
--                   (pure $ C.CDeclarationItem d Nothing Nothing)
--               Nothing ->
--                 pure $ C.CDeclarationItem d Nothing Nothing
--           a -> error (show a)
--       pure (C.CDecl spc decl' ni)
--     C.CDecl spc@[C.CTypeSpec (C.CTypeDef i ni')] decl ni -> do
--       x <- gets (Map.lookup i)
--       case x of
--         Just (Type rst) -> do
--           decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers rst) decl
--           pure $ C.CDecl rst decl' ni
--         Nothing -> do
--           decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers spc) decl
--           pure $ C.CDecl spc decl' ni
--     C.CDecl spc decl ni -> do
--       decl' <- collectNonEmpty' (reduceCDeclarationItem $ identifiers spc) decl
--       pure $ C.CDecl spc decl' ni
--     a -> error (show a)
--    where
--     reduceCDeclarationItem rq' = \case
--       C.CDeclarationItem d i e -> do
--         let (fn, reqs) = cDeclaratorIdentifiers d
--         case fn of
--           Just fn' ->
--             conditionalGivenThat (rq' <> reqs) (Val.is fn')
--           Nothing ->
--             mapM_ (givenThat . Val.is) (rq' <> reqs)
--
--         i' <- optional do
--           liftMaybe i >>= reduce @C.CInitializer
--         e' <- optional do
--           liftMaybe e >>= reduce @C.CExpression
--
--         pure (C.CDeclarationItem d i' e')
--       a -> error (show a)
--
-- instance CReducible C.CInitializer where
--   reduce = \case
--     C.CInitExpr e ni -> reduce @C.CExpression 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' <- collect rmCPartDesignator pds
--       is' <- reduce is
--       pure (pds', is')
--
--     rmCPartDesignator = \case
--       a -> error (show a)
--
-- instance CReducible C.CStatement where
--   reduce = \case
--     C.CCompound is cbi ni -> do
--       cbi' <- collect (reduce @C.CCompoundBlockItem) cbi
--       pure $ C.CCompound is cbi' ni
--     C.CExpr e ni -> do
--       e' <- optional do
--         e' <- liftMaybe e
--         reduce @C.CExpression e'
--       pure $ C.CExpr e' ni
--     C.CIf e s els ni -> do
--       s' <- reduce s
--       e' <- optional do
--         reduce @C.CExpression e
--       els' <- optional do
--         els' <- liftMaybe els
--         given >> reduce els'
--       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
--       reduce s <| do
--         e1' <- reduce @C.CForInit e1
--         e2' <- optional $ liftMaybe e2 >>= reduce @C.CExpression
--         e3' <- optional $ liftMaybe e3 >>= reduce @C.CExpression
--         s' <- reduce s
--         pure $ C.CFor e1' e2' e3' s' ni
--     C.CReturn e ni -> do
--       e' <- traverse (fmap orZero reduce) 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 -> do
--       -- todo fix attrs
--       s' <- reduce s
--       withFallback s' do
--         givenThat (Val.is i)
--         pure $ C.CLabel i s' [] ni
--     C.CGoto i ni ->
--       withFallback (C.CExpr Nothing ni) do
--         givenThat (Val.is i)
--         pure $ C.CGoto i ni
--     C.CWhile e s dow ni -> do
--       e' <- orZero (reduce @C.CExpression e)
--       s' <- reduce s
--       pure $ C.CWhile e' s' dow ni
--     a -> error (show a)
--
-- instance CReducible C.CForInit where
--   reduce = \case
--     C.CForDecl decl -> withFallback (C.CForInitializing Nothing) do
--       C.CForDecl <$> reduce @C.CDeclaration decl
--     C.CForInitializing n -> do
--       C.CForInitializing <$> optional do
--         n' <- liftMaybe n
--         reduce @C.CExpression n'
--
-- instance CReducible C.CExpression where
--   reduce = \case
--     C.CVar i ni -> do
--       givenThat (Val.is i)
--       pure $ C.CVar i ni
--     C.CCall e es ni -> do
--       e' <- reduce e
--       es' <- traverse (fmap orZero reduce) es
--       pure $ C.CCall e' es' ni
--     C.CCond ec et ef ni -> do
--       ec' <- reduce ec
--       ef' <- reduce ef
--       et' <- optional do
--         et' <- liftMaybe et
--         reduce et'
--       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 <- reduce 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' <- reduce @C.CDeclaration cd
--       e' <- reduce 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 -> do
--       e1' <- reduce e1
--       e2' <- orZero (reduce e2)
--       pure $ C.CIndex e1' e2' ni
--     C.CMember e i b ni -> do
--       givenThat (Val.is i)
--       e' <- reduce e
--       pure $ C.CMember e' i b ni
--     C.CComma items ni -> do
--       C.CComma <$> collectNonEmpty' reduce items <*> pure ni
--     e -> error (show e)
--    where
--     onBothExpr elhs erhs = onBoth (reduce elhs) (reduce erhs)
--
-- zeroExp :: C.CExpression C.NodeInfo
-- zeroExp = C.CConst (C.CIntConst (C.cInteger 0) C.undefNode)
--
-- withFallback :: (Alternative m) => a -> m a -> m a
-- withFallback a ma = ma <|> pure a
--
-- orZero :: (Alternative m) => m (C.CExpression C.NodeInfo) -> m (C.CExpression C.NodeInfo)
-- orZero = withFallback zeroExp
--
-- instance CReducible C.CCompoundBlockItem where
--   reduce = \case
--     C.CBlockStmt s ->
--       C.CBlockStmt <$> do
--         given >> reduce @C.CStatement s
--     C.CBlockDecl d ->
--       C.CBlockDecl <$> do
--         reduce @C.CDeclaration d
--     a -> error (show a)