-- Defines the internal state of the application.
module ShellState where



import Control.Monad (liftM, when)
import Data.List (find)
import qualified Data.Map as Map (Map, empty, elems, insert, lookup)
import Data.Maybe (mapMaybe)
import Language.Haskell.FreeTheorems
import Language.Haskell.FreeTheorems.Syntax
import Language.Haskell.FreeTheorems.Theorems
import System.Console.Shell.ShellMonad
import System.IO (hGetBuffering, hSetBuffering, stdout, BufferMode (..))
import Text.PrettyPrint as Doc

import Settings





------- General purpose functions ---------------------------------------------

-- Render pretty documents.
doRender :: Int -> Doc -> String
doRender indent doc = renderStyle (Style PageMode 74 1.2) (nest indent doc)



-- Applies an information retrieval function to the shell state in a shell
-- operation.
fromShellSt :: (ShellState -> a) -> Sh ShellState a
fromShellSt f = liftM f getShellSt





------- Shell state structure -------------------------------------------------


-- The structure of the internal state.
data ShellState = ShellState
  { getParser            :: String -> Parsed [Declaration]
        -- The parser which is used for all inputs.
        -- It should not be changed during the run of the application.
  
  , getDeclarationsFiles :: [FilePath]
        -- The names of all files containing declarations.
  
  , getAllDeclarations   :: Map.Map FilePath [ValidDeclaration]
        -- Maps a file name to the declarations parsed from that file.
  
  , getUserSignatures    :: [ValidSignature]
        -- Contains all signatures entered by the user.
  
  , getCurrentLanguageSubset :: LanguageSubset
        -- The current language subset.
  
  , getCurrentContext    :: Maybe (ValidSignature, Maybe Intermediate)
        -- The current type signature and its interpretation.
        -- The relation variables stored in 'getSpecialisationHistory' are 
        -- specialised in the intermediate.

  , getSpecialisationHistory  :: [Either RelationVariable RelationVariable]
        -- Gives the order of relation variables which were specialised in the
        -- current theorem. The head was specialised first, while the last
        -- relation variables were specialised last.
        -- Left variables are specialised in the standard way, while
        -- right variables are specialised to inverse functions.

  , getCurrentTheoremOptions :: [PrettyTheoremOption]
        -- The current options which modify how theorems are pretty-printed.

  , getSimplifyOutput :: Bool
	-- Whether the Theorems.Simplify module should be used (if allowed by the
	-- language subtype)
  }



-- Initialises the state of the shell from a given settings structure.
-- The declarations and the declarations files are loaded and added afterwards
-- by 'loadDeclarations'.
initialShellState :: Settings -> ShellState
initialShellState settings = ShellState
  { getParser                = parser settings
  , getDeclarationsFiles     = []
  , getAllDeclarations       = Map.empty
  , getUserSignatures        = []
  , getCurrentLanguageSubset = BasicSubset
  , getCurrentContext        = Nothing
  , getSpecialisationHistory = []
  , getCurrentTheoremOptions = [ OmitTypeInstantiations
                               -- , OmitLanguageSubsets
                               ]
  , getSimplifyOutput        = True
  }


-- Resets the state. 
-- All previous state information is lost except the parser and the language
-- subset. Note that also the list of declarations files is erased.
-- These files can be added again by subsequent calls to 'loadDeclarations'.
reset :: ShellState -> ShellState
reset state = ShellState
  { getParser                = getParser state
  , getDeclarationsFiles     = []
  , getAllDeclarations       = Map.empty
  , getUserSignatures        = []
  , getCurrentLanguageSubset = getCurrentLanguageSubset state
  , getCurrentContext        = Nothing
  , getSpecialisationHistory = []
  , getCurrentTheoremOptions = getCurrentTheoremOptions state
  , getSimplifyOutput        = True
  }





------- Loading files ---------------------------------------------------------


-- Loads a declarations file. Returns the modified state, selected error
-- messages and a message containing statistics about loaded declarations.
loadDeclarations :: FilePath -> ShellState -> IO ShellState
loadDeclarations path state = do
  catch
    (do contents <- readFile path
        bufferMode <- hGetBuffering stdout
        hSetBuffering stdout NoBuffering
        putStr $ "Loading `" ++ path ++ "' ... "
        let decls   = concat $ Map.elems $ getAllDeclarations state
            parse   = getParser state
            (vs,es) = runChecks (parse contents >>= checkAgainst decls)

            -- Create a new state by adding the file name and the parsed
            -- declarations.
            state'  = state
                        { getDeclarationsFiles = getDeclarationsFiles state 
                                                 ++ [path]
                        , getAllDeclarations = Map.insert path vs
                                               (getAllDeclarations state) }
            
        if null es
          then  -- Create a message containing statistics.
            putStrLn $ 
              case length vs of
                0         -> "found no declaration."
                1         -> "found 1 declaration."
                otherwise -> "found " ++ show (length vs) ++ " declarations."

          else  -- Create a list of errors.
            putStrLn . doRender 0 $
              case length es of
                1         -> text "found one error:" $$ nest 2 (head es)
                otherwise -> text ("found " ++ show (length es) ++ " errors:")
                                    $$ nest 2 (vcat (take 5 es))

        hSetBuffering stdout bufferMode
        return state'
    )
    (\error -> do
       putStrLn $ "Error: Unable to load the file `" ++ path ++ "'."
       return state 
    )





------- Declaration functions -------------------------------------------------


-- Returns the declarations loaded from the given file.
getDeclarationsIn :: FilePath -> ShellState -> [ValidDeclaration]
getDeclarationsIn file state =
  case Map.lookup file (getAllDeclarations state) of
    Nothing -> []
    Just ds -> ds


-- Returns the names of all loaded declarations and all type signature name 
-- entered by the user.
getAllDeclarationNames :: ShellState -> [String]
getAllDeclarationNames state =
  let allNamesIn f = map (unpackIdent . getDeclarationName . rawDeclaration)
                         (getDeclarationsIn f state)
   in concatMap (\file -> allNamesIn file) (getDeclarationsFiles state)
      ++ getUserSignatureNames state



-- Returns the declaration names loaded from the given file and fitting the
-- given selector function.
getDeclarationNamesIn :: 
    (Declaration -> Bool) -> FilePath -> ShellState -> [String]
getDeclarationNamesIn select file = 
  map (unpackIdent . getDeclarationName) 
  . filter select 
  . map rawDeclaration 
  . getDeclarationsIn file



-- Selector function to retrieve declarations of algebraic data types.
onlyDataDecl :: Declaration -> Bool
onlyDataDecl decl = case decl of
  DataDecl _ -> True
  otherwise  -> False



-- Selector function to retrieve declarations of type renamings.
onlyNewtypeDecl :: Declaration -> Bool
onlyNewtypeDecl decl = case decl of
  NewtypeDecl _ -> True
  otherwise     -> False



-- Selector function to retrieve declarations of type synonyms.
onlyTypeDecl :: Declaration -> Bool
onlyTypeDecl decl = case decl of
  TypeDecl _ -> True
  otherwise  -> False



-- Selector function to retrieve declarations of type classes.
onlyClassDecl :: Declaration -> Bool
onlyClassDecl decl = case decl of
  ClassDecl _ -> True
  otherwise   -> False



-- Selector function to retrieve declarations of type signatures.
onlySigDecl :: Declaration -> Bool
onlySigDecl decl = case decl of
  TypeSig _ -> True
  otherwise -> False



-- Returns the pretty-printed declaration for the given declaration name, if
-- there exists a declaration with such a name.
getDeclarationOf :: String -> ShellState -> Maybe Doc
getDeclarationOf declName state = 
  let us = map (TypeSig . rawSignature) (getUserSignatures state)
      fs = reverse (getDeclarationsFiles state)
      ds = concatMap (\f -> map rawDeclaration (getDeclarationsIn f state)) fs
      getName = unpackIdent . getDeclarationName
   in liftM prettyDeclaration (find (\s -> getName s == declName) (us ++ ds))





------- Signature functions ---------------------------------------------------


-- Returns the names of all type signatures, both loaded from files and entered
-- by the user.
getAllSignatureNames :: ShellState -> [String]
getAllSignatureNames state =
  concatMap (\f -> getDeclarationNamesIn onlySigDecl f state)
            (getDeclarationsFiles state)
  ++ getUserSignatureNames state



-- Returns only the names of type signatures entered by the user.
getUserSignatureNames :: ShellState -> [String]
getUserSignatureNames =
  map (unpackIdent . signatureName . rawSignature) . getUserSignatures



-- Returns all signatures loaded from the given file.
getSignaturesIn :: FilePath -> ShellState -> [ValidSignature]
getSignaturesIn file state =
  case Map.lookup file (getAllDeclarations state) of
    Nothing -> []
    Just ds -> filterSignatures ds



-- Adds a user-entered type signature to the list of known type signatures.
addUserSignature :: ValidSignature -> ShellState -> ShellState
addUserSignature s state = 
  state { getUserSignatures = getUserSignatures state ++ [s] }



-- Returns a type signature for the given name or Nothing, if there is no
-- type signature with that name.
-- Type signatures are searched for beginning from the user-entered type 
-- signatures and then traversing the loaded files in reverse order.
findSignature :: String -> ShellState -> Maybe ValidSignature
findSignature name state = 
  let us = reverse (getUserSignatures state)
      fs = reverse (getDeclarationsFiles state)
      ds = concatMap (\f -> filterSignatures (getDeclarationsIn f state)) fs
      getName = unpackIdent . signatureName . rawSignature
   in find (\s -> getName s == name) (us ++ ds)





------- Context and language functions ----------------------------------------


-- Changes the context to the given type signature.
setCurrentContext :: ValidSignature -> ShellState -> ShellState
setCurrentContext sig state = 
  let l = getCurrentLanguageSubset state
      d = concat . Map.elems . getAllDeclarations $ state
      i = interpret d l sig
   in state { getCurrentContext = Just (sig, i)
            , getSpecialisationHistory = [] }



-- Changes the current language subset.
setCurrentLanguageSubset :: LanguageSubset -> ShellState -> ShellState
setCurrentLanguageSubset l state =
  case getCurrentContext state of
    Nothing    -> state { getCurrentLanguageSubset = l }
    Just (s,_) -> 
      let h = getSpecialisationHistory state
          d = concat . Map.elems . getAllDeclarations $ state
          f i = either (specialise i) (specialiseInverse i)
          i = fmap (\i -> foldl f i h) (interpret d l s)
       in state { getCurrentContext = Just (s,i)
                , getCurrentLanguageSubset = l }



-- Adds a relation variable to the end of the specialisation history and 
-- specialises the given relation variable in the current theorem.
addSpecialisation :: 
    Either RelationVariable RelationVariable -> ShellState -> ShellState
addSpecialisation rv state =
  case getCurrentContext state of
    Nothing    -> state
    Just (s,i) ->
      let h  = getSpecialisationHistory state ++ [rv]
          i' = fmap (\i -> either (specialise i) (specialiseInverse i) rv) i
       in state { getCurrentContext = Just (s, i')
                , getSpecialisationHistory = h }



-- Reverts the last specialisation, if there is any.
undoSpecialisation :: ShellState -> ShellState
undoSpecialisation state = 
  if null (getSpecialisationHistory state)
    then state
    else case getCurrentContext state of
           Nothing    -> state
           Just (s,_) -> let h = init (getSpecialisationHistory state)
                             l = getCurrentLanguageSubset state
                             d = concat . Map.elems . getAllDeclarations $ state
                             f i = either (specialise i) (specialiseInverse i)
                             i = fmap (\i -> foldl f i h) (interpret d l s)
                          in state { getSpecialisationHistory = h
                                   , getCurrentContext = Just (s, i) }