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

module ReduceC where

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)

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