{-# LANGUAGE CPP, PatternGuards #-}
module Agda.TypeChecking.Monad.Signature where

import Control.Monad.State
import Control.Monad.Reader
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map (Map)
import qualified Data.Map as Map
import Data.List
import Data.Function
import qualified Agda.Utils.IO.Locale as LocIO

import Agda.Syntax.Abstract.Name
import Agda.Syntax.Common
import Agda.Syntax.Internal
import Agda.Syntax.Position

import Agda.TypeChecking.Monad.Base
import Agda.TypeChecking.Monad.Context
import Agda.TypeChecking.Monad.Options
import Agda.TypeChecking.Monad.Env
import Agda.TypeChecking.Monad.Mutual
import Agda.TypeChecking.Monad.Open
import Agda.TypeChecking.Substitute
import {-# SOURCE #-} Agda.TypeChecking.Polarity

import Agda.Utils.Monad
import Agda.Utils.Map as Map
import Agda.Utils.Size
import Agda.Utils.Permutation

#include "../../undefined.h"
import Agda.Utils.Impossible

modifySignature :: MonadTCM tcm => (Signature -> Signature) -> tcm ()
modifySignature f = modify $ \s -> s { stSignature = f $ stSignature s }

modifyImportedSignature :: MonadTCM tcm => (Signature -> Signature) -> tcm ()
modifyImportedSignature f = modify $ \s -> s { stImports = f $ stImports s }

getSignature :: MonadTCM tcm => tcm Signature
getSignature = liftTCM $ gets stSignature

getImportedSignature :: MonadTCM tcm => tcm Signature
getImportedSignature = liftTCM $ gets stImports

setSignature :: MonadTCM tcm => Signature -> tcm ()
setSignature sig = modifySignature $ const sig

setImportedSignature :: MonadTCM tcm => Signature -> tcm ()
setImportedSignature sig = liftTCM $ modify $ \s -> s { stImports = sig }

withSignature :: MonadTCM tcm => Signature -> tcm a -> tcm a
withSignature sig m =
    do	sig0 <- getSignature
	setSignature sig
	r <- m
	setSignature sig0
        return r

-- | Add a constant to the signature. Lifts the definition to top level.
addConstant :: MonadTCM tcm => QName -> Definition -> tcm ()
addConstant q d = liftTCM $ do
  tel <- getContextTelescope
  let tel' = killRange $ case theDef d of
	      Constructor{} -> hideTel tel
	      _		    -> tel
  modifySignature $ \sig -> sig
    { sigDefinitions = Map.insertWith (+++) q (abstract tel' d') $ sigDefinitions sig }
  i <- currentMutualBlock
  setMutualBlock i q
  where
    d' = d { defName = q }
    new +++ old = new { defDisplay = defDisplay new ++ defDisplay old }

    hideTel  EmptyTel		      = EmptyTel
    hideTel (ExtendTel (Arg _ t) tel) = ExtendTel (Arg Hidden t) $ hideTel <$> tel

addHaskellCode :: MonadTCM tcm => QName -> HaskellType -> HaskellCode -> tcm ()
addHaskellCode q hsTy hsDef =
  -- TODO: sanity checking
  modifySignature $ \sig -> sig
  { sigDefinitions = Map.adjust addHs q $ sigDefinitions sig }
  where
    addHs def@Defn{theDef = con@Constructor{}} =
      def{theDef = con{conHsCode = Just (hsTy, hsDef)}}
    addHs def@Defn{theDef = ax@Axiom{}} =
      def{theDef = ax{axHsDef = Just $ HsDefn hsTy hsDef}}
    addHs def = def

addHaskellType :: MonadTCM tcm => QName -> HaskellType -> tcm ()
addHaskellType q hsTy =
  -- TODO: sanity checking
  modifySignature $ \sig -> sig
  { sigDefinitions = Map.adjust addHs q $ sigDefinitions sig }
  where
    addHs def@Defn{theDef = ax@Axiom{}} =
      def{theDef = ax{axHsDef = Just $ HsType hsTy}}
    addHs def@Defn{theDef = d@Datatype{}} =
      def{theDef = d{dataHsType = Just hsTy}}
    addHs def = def

unionSignatures :: [Signature] -> Signature
unionSignatures ss = foldr unionSignature emptySignature ss
  where
    unionSignature (Sig a b) (Sig c d) = Sig (Map.union a c) (Map.union b d)

-- | Add a section to the signature.
addSection :: MonadTCM tcm => ModuleName -> Nat -> tcm ()
addSection m fv = do
  tel <- getContextTelescope
  let sec = Section tel fv
  modifySignature $ \sig -> sig { sigSections = Map.insert m sec $ sigSections sig }

-- | Lookup a section. If it doesn't exist that just means that the module
--   wasn't parameterised.
lookupSection :: MonadTCM tcm => ModuleName -> tcm Telescope
lookupSection m = do
  sig  <- sigSections <$> getSignature
  isig <- sigSections <$> getImportedSignature
  return $ maybe EmptyTel secTelescope $ Map.lookup m sig `mplus` Map.lookup m isig

-- Add display forms to all names @xn@ such that @x = x1 es1@, ... @xn-1 = xn esn@.
addDisplayForms :: QName -> TCM ()
addDisplayForms x = do
  args <- getContextArgs
  add args x x []
  where
    add args top x ps = do
      cs <- defClauses <$> getConstInfo x
      case cs of
	[ Clause{ clauseBody = b } ]
	  | Just (m, Def y vs) <- strip b
	  , m == length args && args `isPrefixOf` vs -> do
	      let ps' = raise 1 (map unArg vs) ++ ps
	      reportSLn "tc.section.apply.display" 20 $ "adding display form " ++ show y ++ " --> " ++ show top
	      addDisplayForm y (Display 0 ps' $ DTerm $ Def top args)
	      add args top y $ drop (length args) ps'
	_ -> do
	      let reason = case cs of
		    []    -> "no clauses"
		    _:_:_ -> "many clauses"
		    [ Clause{ clauseBody = b } ] -> case strip b of
		      Nothing -> "bad body"
		      Just (m, Def y vs)
			| m < length args -> "too few args"
			| m > length args -> "too many args"
			| otherwise	      -> "args=" ++ unwords (map var args) ++ " vs=" ++ unwords (map var vs)
			  where
			    var (Arg h x) = hid h $ case x of
			      Var i []	-> show i
			      MetaV _ _	-> "?"
			      _		-> "_"
			    hid NotHidden s = s
			    hid Hidden    s = "{" ++ s ++ "}"
		      Just (m, v) -> "not a def body"
	      reportSLn "tc.section.apply.display" 30 $ "no display form from" ++ show x ++ " because " ++ reason
	      return ()
    strip (Body v)   = return (0, v)
    strip  NoBody    = Nothing
    strip (NoBind b) = Nothing
    strip (Bind b)   = do
      (n, v) <- strip $ absBody b
      return (n + 1, v)

applySection ::
  MonadTCM tcm => ModuleName -> Telescope -> ModuleName -> Args ->
  Map QName QName -> Map ModuleName ModuleName -> tcm ()
applySection new ptel old ts rd rm = liftTCM $ do
  sig  <- getSignature
  isig <- getImportedSignature
  let ss = Map.toList $ Map.filterKeys partOfOldM $ sigSections sig `Map.union` sigSections isig
      ds = Map.toList $ Map.filterKeys partOfOldD $ sigDefinitions sig `Map.union` sigDefinitions isig
  mapM_ (copyDef ts) ds
  mapM_ (copySec ts) ss
  where
    partOfOldM x = x `isSubModuleOf` old
    partOfOldD x = x `isInModule`    old

    copyName x = maybe x id $ Map.lookup x rd

    copyDef :: Args -> (QName, Definition) -> TCM ()
    copyDef ts (x, d) = case Map.lookup x rd of
	Nothing -> return ()  -- if it's not in the renaming it was private and
			      -- we won't need it
	Just y	-> do
	  addConstant y (nd y)
          computePolarity y
	  -- Set display form for the old name if it's not a constructor.
	  unless (isCon || size ptel > 0) $ do
	    addDisplayForms y
      where
	t  = defType d `apply` ts
	-- the name is set by the addConstant function
	nd y = Defn y t [] (-1) def  -- TODO: mutual block?
        oldDef = theDef d
	isCon = case oldDef of
	  Constructor{} -> True
	  _		-> False
        -- TODO: compute polarity for the new definition
	def  = case oldDef of
                Constructor{ conPars = np, conData = d } ->
                  oldDef { conPars = np - size ts, conData = copyName d }
                Datatype{ dataPars = np, dataCons = cs } ->
                  oldDef { dataPars = np - size ts, dataClause = Just cl, dataCons = map copyName cs }
                Record{ recPars = np, recConType = t, recTel = tel } ->
                  oldDef { recPars = np - size ts, recClause = Just cl
                         , recConType = apply t ts, recTel = apply tel ts
                         }
		_ ->
                  Function { funClauses        = [cl]
                           , funDelayed        = NotDelayed
                           , funInv            = NotInjective
                           , funPolarity       = []
                           , funArgOccurrences = []
                           , funAbstr          = ConcreteDef
                           }
	cl = Clause { clauseRange = getRange $ defClauses d
                    , clauseTel   = EmptyTel
                    , clausePerm  = idP 0
                    , clausePats  = []
                    , clauseBody  = Body $ Def x ts
                    }

    copySec :: Args -> (ModuleName, Section) -> TCM ()
    copySec ts (x, sec) = case Map.lookup x rm of
	Nothing -> return ()  -- if it's not in the renaming it was private and
			      -- we won't need it
	Just y  -> addCtxTel (apply tel ts) $ addSection y 0
      where
	tel = secTelescope sec

addDisplayForm :: MonadTCM tcm => QName -> DisplayForm -> tcm ()
addDisplayForm x df = do
  d <- makeOpen df
  modifyImportedSignature (add d)
  modifySignature (add d)
  where
    add df sig = sig { sigDefinitions = Map.adjust addDf x defs }
      where
	addDf def = def { defDisplay = df : defDisplay def }
	defs	  = sigDefinitions sig

canonicalName :: MonadTCM tcm => QName -> tcm QName
canonicalName x = do
  def <- theDef <$> getConstInfo x
  case def of
    Constructor{conSrcCon = c}                                -> return c
    Record{recClause = Just (Clause{ clauseBody = body })}    -> canonicalName $ extract body
    Datatype{dataClause = Just (Clause{ clauseBody = body })} -> canonicalName $ extract body
    _                                                         -> return x
  where
    extract NoBody	     = __IMPOSSIBLE__
    extract (Body (Def x _)) = x
    extract (Body _)	     = __IMPOSSIBLE__
    extract (Bind (Abs _ b)) = extract b
    extract (NoBind b)	     = extract b

-- | Can be called on either a (co)datatype, a record type or a
--   (co)constructor.
whatInduction :: MonadTCM tcm => QName -> tcm Induction
whatInduction c = do
  def <- theDef <$> getConstInfo c
  case def of
    Datatype{ dataInduction = i } -> return i
    Record{}                      -> return Inductive
    Constructor{ conInd = i }     -> return i
    _                             -> __IMPOSSIBLE__

-- | Does the given constructor come from a single-constructor type?
--
-- Precondition: The name has to refer to a constructor.
singleConstructorType :: QName -> TCM Bool
singleConstructorType q = do
  d <- theDef <$> getConstInfo q
  case d of
    Record {}                   -> return True
    Constructor { conData = d } -> do
      di <- theDef <$> getConstInfo d
      return $ case di of
        Record {}                  -> True
        Datatype { dataCons = cs } -> length cs == 1
        _                          -> __IMPOSSIBLE__
    _ -> __IMPOSSIBLE__

-- | Lookup the definition of a name. The result is a closed thing, all free
--   variables have been abstracted over.
getConstInfo :: MonadTCM tcm => QName -> tcm Definition
getConstInfo q = liftTCM $ join $ pureTCM $ \st env ->
  let defs  = sigDefinitions $ stSignature st
      idefs = sigDefinitions $ stImports st
      ab    = treatAbstractly' q env
      smash = (++) `on` maybe [] (:[])
  in case smash (Map.lookup q defs) (Map.lookup q idefs) of
      []  -> fail $ "Unbound name: " ++ show q
      [d] -> mkAbs ab d
      ds  -> fail $ "Ambiguous name: " ++ show q
  where
    mkAbs True d =
      case makeAbstract d of
	Just d	-> return d
	Nothing	-> typeError $ NotInScope [qnameToConcrete q]
	  -- the above can happen since the scope checker is a bit sloppy with 'abstract'
    mkAbs False d = return d

-- | Look up the polarity of a definition.
getPolarity :: MonadTCM tcm => QName -> tcm [Polarity]
getPolarity q = liftTCM $ do
  defn <- theDef <$> getConstInfo q
  case defn of
    Function{ funPolarity  = p } -> return p
    Datatype{ dataPolarity = p } -> return p
    Record{ recPolarity    = p } -> return p
    _                            -> return []

getPolarity' :: MonadTCM tcm => Comparison -> QName -> tcm [Polarity]
getPolarity' CmpEq _  = return []
getPolarity' CmpLeq q = getPolarity q

-- | Set the polarity of a definition.
setPolarity :: MonadTCM tcm => QName -> [Polarity] -> tcm ()
setPolarity q pol = liftTCM $ do
  modifySignature setP
  where
    setP sig = sig { sigDefinitions = Map.adjust setPx q defs }
      where
	setPx def = def { theDef = setPd $ theDef def }
        setPd d   = case d of
          Function{} -> d { funPolarity  = pol }
          Datatype{} -> d { dataPolarity = pol }
          Record{}   -> d { recPolarity  = pol }
          _          -> d
	defs	  = sigDefinitions sig

getArgOccurrence :: MonadTCM tcm => QName -> Nat -> tcm Occurrence
getArgOccurrence d i = do
  def <- theDef <$> getConstInfo d
  return $ case def of
    Function { funArgOccurrences  = os } -> look i os
    Datatype { dataArgOccurrences = os } -> look i os
    Record   { recArgOccurrences  = os } -> look i os
    Constructor{}                        -> Positive
    _                                    -> Negative
  where
    look i os = (os ++ repeat Negative) !! fromIntegral i

setArgOccurrences :: MonadTCM tcm => QName -> [Occurrence] -> tcm ()
setArgOccurrences d os = liftTCM $
  modifySignature setO
  where
    setO sig = sig { sigDefinitions = Map.adjust setOx d defs }
      where
	setOx def = def { theDef = setOd $ theDef def }
        setOd d   = case d of
          Function{} -> d { funArgOccurrences  = os }
          Datatype{} -> d { dataArgOccurrences = os }
          Record{}   -> d { recArgOccurrences  = os }
          _          -> d
	defs	  = sigDefinitions sig


-- | Look up the number of free variables of a section. This is equal to the
--   number of parameters if we're currently inside the section and 0 otherwise.
getSecFreeVars :: MonadTCM tcm => ModuleName -> tcm Nat
getSecFreeVars m = do
  sig  <- sigSections <$> getSignature
  isig <- sigSections <$> getImportedSignature
  top <- currentModule
  case top `isSubModuleOf` m || top == m of
    True  -> return $ maybe 0 secFreeVars $ Map.lookup m (Map.union sig isig)
    False -> return 0

-- | Compute the number of free variables of a module. This is the sum of
--   the free variables of its sections.
getModuleFreeVars :: MonadTCM tcm => ModuleName -> tcm Nat
getModuleFreeVars m = sum <$> ((:) <$> getAnonymousVariables m <*> mapM getSecFreeVars ms)
  where
    ms = map mnameFromList . inits . mnameToList $ m

-- | Compute the number of free variables of a defined name. This is the sum of
--   the free variables of the sections it's contained in.
getDefFreeVars :: MonadTCM tcm => QName -> tcm Nat
getDefFreeVars q = getModuleFreeVars (qnameModule q)

-- | Compute the context variables to apply a definition to.
freeVarsToApply :: MonadTCM tcm => QName -> tcm Args
freeVarsToApply x = genericTake <$> getDefFreeVars x <*> getContextArgs

-- | Instantiate a closed definition with the correct part of the current
--   context.
instantiateDef :: MonadTCM tcm => Definition -> tcm Definition
instantiateDef d = do
  vs  <- freeVarsToApply $ defName d
  verboseS "tc.sig.inst" 30 $ do
    ctx <- getContext
    m   <- currentModule
    liftIO $ LocIO.putStrLn $ "instDef in " ++ show m ++ ": " ++ show (defName d) ++ " " ++
			unwords (map show . take (size vs) . reverse . map (fst . unArg) $ ctx)
  return $ d `apply` vs

-- | Give the abstract view of a definition.
makeAbstract :: Definition -> Maybe Definition
makeAbstract d = do def <- makeAbs $ theDef d
		    return d { theDef = def }
    where
	makeAbs Datatype   {dataAbstr = AbstractDef} = Just $ Axiom Nothing
	makeAbs Function   {funAbstr  = AbstractDef} = Just $ Axiom Nothing
	makeAbs Constructor{conAbstr  = AbstractDef} = Nothing
	makeAbs d                                    = Just d

-- | Enter abstract mode. Abstract definition in the current module are transparent.
inAbstractMode :: MonadTCM tcm => tcm a -> tcm a
inAbstractMode = local $ \e -> e { envAbstractMode = AbstractMode }

-- | Not in abstract mode. All abstract definitions are opaque.
inConcreteMode :: MonadTCM tcm => tcm a -> tcm a
inConcreteMode = local $ \e -> e { envAbstractMode = ConcreteMode }

-- | Ignore abstract mode. All abstract definitions are transparent.
ignoreAbstractMode :: MonadTCM tcm => tcm a -> tcm a
ignoreAbstractMode = local $ \e -> e { envAbstractMode = IgnoreAbstractMode }

-- | Check whether a name might have to be treated abstractly (either if we're
--   'inAbstractMode' or it's not a local name). Returns true for things not
--   declared abstract as well, but for those 'makeAbstract' will have no effect.
treatAbstractly :: MonadTCM tcm => QName -> tcm Bool
treatAbstractly q = treatAbstractly' q <$> ask

treatAbstractly' :: QName -> TCEnv -> Bool
treatAbstractly' q env = case envAbstractMode env of
  ConcreteMode	     -> True
  IgnoreAbstractMode -> False
  AbstractMode	     -> not $ current == m || current `isSubModuleOf` m
  where
    current = envCurrentModule env
    m	    = qnameModule q

-- | get type of a constant
typeOfConst :: MonadTCM tcm => QName -> tcm Type
typeOfConst q = defType <$> (instantiateDef =<< getConstInfo q)

-- | The name must be a datatype.
sortOfConst :: MonadTCM tcm => QName -> tcm Sort
sortOfConst q =
    do	d <- theDef <$> getConstInfo q
	case d of
	    Datatype{dataSort = s} -> return s
	    _			   -> fail $ "Expected " ++ show q ++ " to be a datatype."