-- |
-- Type classes and built-in implementation for primitive Haskell types
-- 

module Z3.Base.Class (
    -- ** Types whose values are encodable to Z3 internal AST
    Z3Encoded(..),
    -- ** Types representable as Z3 Sort
    -- XXX: Unsound now
    -- XXX: Too flexible, can be used to encode Type ADT
    Z3Sorted(..),
    -- ** Type proxy helper, used with Z3Sorted
    Z3Sort(..),
    -- ** Types with reserved value for Z3 encoding use
    -- XXX: Magic value for built-in types
    Z3Reserved(..),
    -- ** Monad which can be instantiated into a concrete context
    SMT(..),
    Datatypes,
    ValBindings,
    emptyValBindings,
    -- ** Heterogenous list, a hack to encode different "term" into a list
    -- Used to encode function argument list
    HeteroList(..),
    mapH
) where

import Z3.Monad

import Control.Monad.Except

import qualified Data.Map as M
import qualified Data.Set as S

data Z3Sort a = Z3Sort

class Z3Encoded a where
    encode :: SMT m e => a -> m e AST

type Datatypes ty = [(String, [(String, [(String, ty)])])]
type ValBindings = ([String], HeteroList)

emptyValBindings :: ValBindings
emptyValBindings = ([], Nil)


-- | XXX: Unsound
class Z3Sorted a where
    -- | Map a value to Sort, the value should be a type-level thing
    sortOf :: SMT m e => a -> m e Sort
    sortOf _ = sortPhantom (Z3Sort :: Z3Sort a)

    -- | Map a Haskell type to Sort
    sortPhantom :: SMT m e => Z3Sort a -> m e Sort
    sortPhantom _ = smtError "sort error"

class Z3Encoded a => Z3Reserved a where
    def :: a

class (MonadError String (m e), MonadZ3 (m e)) => SMT m e where
    -- | Globally unique id
    genFreshId :: m e Int

    -- | Given data type declarations, extra field, and the SMT monad, return the fallible result in IO monad
    runSMT :: Z3Sorted ty => ValBindings -> Datatypes ty -> e -> m e a -> IO (Either String a)

    bindVals :: ([String], HeteroList) -> m e ()
    bindVals (x:xs, Cons v vs) = do
        ast <- encode v
        st <- sortOf v
        bindVal x ast st
        bindVals (xs, vs)
    bindVals _ = return ()

    modifyValBindCtx :: ValBindings -> m e ()

    -- | Binding a variable String name to two things: an de Brujin idx as Z3 AST generated by mkBound and binder's Sort
    bindVal :: String -> AST -> Sort -> m e ()

    getValBindMaybe :: String -> m e (Maybe (AST, Sort))

    getValBind :: String -> m e (AST, Sort)
    getValBind x = getValBindMaybe x >>= \case
        Just p  -> return p
        Nothing -> smtError $ "no such variable: " ++ x

    getDataTypeMaybe :: String -> m e (Maybe Sort)

    getDataType :: String -> m e Sort
    getDataType x = getDataTypeMaybe x >>= \case
        Just p  -> return p
        Nothing -> smtError $ "no such type variable: " ++ x

    -- | Get extra
    getExtra :: m e e

    -- | Set extra
    modifyExtra :: (e -> e) -> m e ()

    -- | User don't have to import throwError
    smtError :: String -> m e a
    smtError = throwError

instance Z3Reserved Int where
    def = -1 -- XXX: Magic number

instance Z3Sorted Int where
    sortPhantom _ = mkIntSort

instance Z3Encoded Int where
    encode i = mkIntSort >>= mkInt i

instance Z3Reserved Double where
    def = -1.0 -- XXX: Magic number

instance Z3Sorted Double where
    sortPhantom _ = mkRealSort

instance Z3Encoded Double where
    encode = mkRealNum

instance Z3Reserved Bool where
    def = False -- XXX: Magic number

instance Z3Sorted Bool where
    sortPhantom _ = mkBoolSort

instance Z3Encoded Bool where
    encode = mkBool

-- The basic idea:
-- For each (k, v), assert in Z3 that if we select k from array we will get
-- the same value v
-- HACK: to set a default value for rest fields (or else we always get the last asserted value
--       as default, which is certainly not complying to finite map's definition), thus the
--       user should guarantee that he/she will never never think this value as a vaid one,
--       if not, he/she might get "a valid value mapped to a invalid key" semantics
instance (Z3Sorted k, Z3Encoded k, Z3Sorted v, Z3Reserved v) => Z3Encoded (M.Map k v) where
    encode m = do
        fid <- genFreshId
        arrSort <- sortOf m
        arr <- mkFreshConst ("map" ++ "_" ++ show fid) arrSort
        mapM_ (\(k, v) -> do
            kast <- encode k
            vast <- encode v
            sel <- mkSelect arr kast
            mkEq sel vast >>= assert) (M.toList m)
        arrValueDef <- mkArrayDefault arr
        vdef <- encode (def :: v)
        mkEq arrValueDef vdef >>= assert
        return arr

instance (Z3Sorted k, Z3Sorted v) => Z3Sorted (M.Map k v) where
    sortPhantom _ = do
        sk <- sortPhantom  (Z3Sort :: Z3Sort k)
        sv <- sortPhantom  (Z3Sort :: Z3Sort v)
        mkArraySort sk sv

-- Basic idea:
-- Set v =def= Map v {0, 1}
-- Thank god, this is much more sound
instance (Z3Sorted v, Z3Encoded v) => Z3Encoded (S.Set v) where
    encode s = do
        setSort <- sortOf s
        fid <- genFreshId
        arr <- mkFreshConst ("set" ++ "_" ++ show fid) setSort
        mapM_ (\e -> do
            ast <- encode e
            sel <- mkSelect arr ast
            one <- (mkIntSort >>= mkInt 1)
            mkEq sel one >>= assert) (S.toList s)
        arrValueDef <- mkArrayDefault arr
        zero <- (mkIntSort >>= mkInt 0)
        mkEq zero arrValueDef >>= assert
        return arr

instance Z3Sorted v => Z3Sorted (S.Set v) where
    sortPhantom _ = do
        sortElem <- sortPhantom (Z3Sort :: Z3Sort v)
        intSort <- mkIntSort
        mkArraySort sortElem intSort


data HeteroList where
    Cons :: forall a. (Z3Sorted a, Z3Encoded a) => a -> HeteroList -> HeteroList
    Nil :: HeteroList

instance Eq HeteroList where
  Nil == Nil = True
  Cons _ h1 == Cons _ h2 = h1 == h2
  _ == _ = False

mapH :: (forall a. (Z3Sorted a, Z3Encoded a) => a -> b) -> HeteroList -> [b]
mapH _ Nil = []
mapH f (Cons a l) = f a : mapH f l