module Agda.Syntax.Concrete.Definitions.Monad where

import Control.Monad.Except
import Control.Monad.State

import Data.Bifunctor (second)
import Data.Map (Map)
import qualified Data.Map as Map

import Agda.Syntax.Position
import Agda.Syntax.Common hiding (TerminationCheck())
import Agda.Syntax.Concrete.Name
import Agda.Syntax.Concrete.Definitions.Types
import Agda.Syntax.Concrete.Definitions.Errors

import Agda.Utils.CallStack ( CallStack, HasCallStack, withCallerCallStack )
import Agda.Utils.Lens

import Agda.Utils.Impossible

-- | Nicifier monad.
--   Preserve the state when throwing an exception.

newtype Nice a = Nice { unNice :: ExceptT DeclarationException (State NiceEnv) a }
  deriving ( Functor, Applicative, Monad
           , MonadState NiceEnv, MonadError DeclarationException
           )

-- | Run a Nicifier computation, return result and warnings
--   (in chronological order).
runNice :: Nice a -> (Either DeclarationException a, NiceWarnings)
runNice m = second (reverse . niceWarn) $
  runExceptT (unNice m) `runState` initNiceEnv

-- | Nicifier state.

data NiceEnv = NiceEnv
  { _loneSigs :: LoneSigs
    -- ^ Lone type signatures that wait for their definition.
  , _termChk  :: TerminationCheck
    -- ^ Termination checking pragma waiting for a definition.
  , _posChk   :: PositivityCheck
    -- ^ Positivity checking pragma waiting for a definition.
  , _uniChk   :: UniverseCheck
    -- ^ Universe checking pragma waiting for a data/rec signature or definition.
  , _catchall :: Catchall
    -- ^ Catchall pragma waiting for a function clause.
  , _covChk  :: CoverageCheck
    -- ^ Coverage pragma waiting for a definition.
  , niceWarn :: NiceWarnings
    -- ^ Stack of warnings. Head is last warning.
  , _nameId  :: NameId
    -- ^ We distinguish different 'NoName's (anonymous definitions) by a unique 'NameId'.
  }

data LoneSig = LoneSig
  { loneSigRange :: Range
  , loneSigName  :: Name
      -- ^ If 'isNoName', this name can have a different 'NameId'
      --   than the key of 'LoneSigs' pointing to it.
  , loneSigKind  :: DataRecOrFun
  }

type LoneSigs     = Map Name LoneSig
     -- ^ We retain the 'Name' also in the codomain since
     --   'Name' as a key is up to @Eq Name@ which ignores the range.
     --   However, without range names are not unique in case the
     --   user gives a second definition of the same name.
     --   This causes then problems in 'replaceSigs' which might
     --   replace the wrong signature.
     --
     --   Another reason is that we want to distinguish different
     --   occurrences of 'NoName' in a mutual block (issue #4157).
     --   The 'NoName' in the codomain will have a unique 'NameId'.

type NiceWarnings = [DeclarationWarning]
     -- ^ Stack of warnings. Head is last warning.

-- | Initial nicifier state.

initNiceEnv :: NiceEnv
initNiceEnv = NiceEnv
  { _loneSigs = Map.empty
  , _termChk  = TerminationCheck
  , _posChk   = YesPositivityCheck
  , _uniChk   = YesUniverseCheck
  , _catchall = False
  , _covChk   = YesCoverageCheck
  , niceWarn  = []
  , _nameId   = NameId 1 noModuleNameHash
  }

lensNameId :: Lens' NameId NiceEnv
lensNameId f e = f (_nameId e) <&> \ i -> e { _nameId = i }

nextNameId :: Nice NameId
nextNameId = do
  i <- use lensNameId
  lensNameId %= succ
  return i

-- * Handling the lone signatures, stored to infer mutual blocks.

-- | Lens for field '_loneSigs'.

loneSigs :: Lens' LoneSigs NiceEnv
loneSigs f e = f (_loneSigs e) <&> \ s -> e { _loneSigs = s }

-- | Adding a lone signature to the state.
--   Return the name (which is made unique if 'isNoName').

addLoneSig :: Range -> Name -> DataRecOrFun -> Nice Name
addLoneSig r x k = do
  -- Andreas, 2020-05-19, issue #4157, make '_' unique.
  x' <- if not $ isNoName x then return x else do
    i <- nextNameId
    return x{ nameId = i }
  loneSigs %== \ s -> do
    let (mr, s') = Map.insertLookupWithKey (\ _k new _old -> new) x (LoneSig r x' k) s
    case mr of
      Nothing -> return s'
      Just{}  -> declarationException $
        if not $ isNoName x then DuplicateDefinition x else DuplicateAnonDeclaration r
  return x'

-- | Remove a lone signature from the state.

removeLoneSig :: Name -> Nice ()
removeLoneSig x = loneSigs %= Map.delete x

-- | Search for forward type signature.

getSig :: Name -> Nice (Maybe DataRecOrFun)
getSig x = fmap loneSigKind . Map.lookup x <$> use loneSigs

-- | Check that no lone signatures are left in the state.

noLoneSigs :: Nice Bool
noLoneSigs = null <$> use loneSigs

-- | Ensure that all forward declarations have been given a definition.

forgetLoneSigs :: Nice ()
forgetLoneSigs = loneSigs .= Map.empty

checkLoneSigs :: LoneSigs -> Nice ()
checkLoneSigs xs = do
  forgetLoneSigs
  unless (Map.null xs) $ declarationWarning $ MissingDefinitions $
    map (\s -> (loneSigName s , loneSigRange s)) $ Map.elems xs

-- | Get names of lone function signatures, plus their unique names.

loneFuns :: LoneSigs -> [(Name,Name)]
loneFuns = map (second loneSigName) . filter (isFunName . loneSigKind . snd) . Map.toList

-- | Create a 'LoneSigs' map from an association list.

loneSigsFromLoneNames :: [(Range, Name, DataRecOrFun)] -> LoneSigs
loneSigsFromLoneNames = Map.fromListWith __IMPOSSIBLE__ . map (\(r,x,k) -> (x, LoneSig r x k))

-- | Lens for field '_termChk'.

terminationCheckPragma :: Lens' TerminationCheck NiceEnv
terminationCheckPragma f e = f (_termChk e) <&> \ s -> e { _termChk = s }

withTerminationCheckPragma :: TerminationCheck -> Nice a -> Nice a
withTerminationCheckPragma tc f = do
  tc_old <- use terminationCheckPragma
  terminationCheckPragma .= tc
  result <- f
  terminationCheckPragma .= tc_old
  return result

coverageCheckPragma :: Lens' CoverageCheck NiceEnv
coverageCheckPragma f e = f (_covChk e) <&> \ s -> e { _covChk = s }

withCoverageCheckPragma :: CoverageCheck -> Nice a -> Nice a
withCoverageCheckPragma tc f = do
  tc_old <- use coverageCheckPragma
  coverageCheckPragma .= tc
  result <- f
  coverageCheckPragma .= tc_old
  return result

-- | Lens for field '_posChk'.

positivityCheckPragma :: Lens' PositivityCheck NiceEnv
positivityCheckPragma f e = f (_posChk e) <&> \ s -> e { _posChk = s }

withPositivityCheckPragma :: PositivityCheck -> Nice a -> Nice a
withPositivityCheckPragma pc f = do
  pc_old <- use positivityCheckPragma
  positivityCheckPragma .= pc
  result <- f
  positivityCheckPragma .= pc_old
  return result

-- | Lens for field '_uniChk'.

universeCheckPragma :: Lens' UniverseCheck NiceEnv
universeCheckPragma f e = f (_uniChk e) <&> \ s -> e { _uniChk = s }

withUniverseCheckPragma :: UniverseCheck -> Nice a -> Nice a
withUniverseCheckPragma uc f = do
  uc_old <- use universeCheckPragma
  universeCheckPragma .= uc
  result <- f
  universeCheckPragma .= uc_old
  return result

-- | Get universe check pragma from a data/rec signature.
--   Defaults to 'YesUniverseCheck'.

getUniverseCheckFromSig :: Name -> Nice UniverseCheck
getUniverseCheckFromSig x = maybe YesUniverseCheck universeCheck <$> getSig x

-- | Lens for field '_catchall'.

catchallPragma :: Lens' Catchall NiceEnv
catchallPragma f e = f (_catchall e) <&> \ s -> e { _catchall = s }

-- | Get current catchall pragma, and reset it for the next clause.

popCatchallPragma :: Nice Catchall
popCatchallPragma = do
  ca <- use catchallPragma
  catchallPragma .= False
  return ca

withCatchallPragma :: Catchall -> Nice a -> Nice a
withCatchallPragma ca f = do
  ca_old <- use catchallPragma
  catchallPragma .= ca
  result <- f
  catchallPragma .= ca_old
  return result

-- | Add a new warning.
niceWarning :: DeclarationWarning -> Nice ()
niceWarning w = modify $ \ st -> st { niceWarn = w : niceWarn st }

declarationException :: HasCallStack => DeclarationException' -> Nice a
declarationException e = withCallerCallStack $ throwError . flip DeclarationException e

declarationWarning' :: DeclarationWarning' -> CallStack -> Nice ()
declarationWarning' w loc = niceWarning $ DeclarationWarning loc w

declarationWarning :: HasCallStack => DeclarationWarning' -> Nice ()
declarationWarning = withCallerCallStack . declarationWarning'