{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}

module Task2
  ( MonadJS(..)
  , VarJS(..)
  , runSTJS
  , runSTJS1
  , f1
  , f2
  , f3
  ) where

import           Control.Applicative (liftA2)
import           Control.Monad       (when)
import           Control.Monad.ST    (ST, runST)
import           Data.Semigroup      (Semigroup, (<>))
import           Data.STRef          (STRef, newSTRef, readSTRef, writeSTRef)

data Number
  = Inter Int
  | Frac Double
  deriving (Show)

instance Semigroup Number where
  Inter x <> Inter y = Inter $ x + y
  Frac x <> Frac y = Frac $ x + y
  Frac x <> Inter y = Frac $ x + fromIntegral y
  a <> b = b <> a

instance Eq Number where
  Inter x == Inter y = x == y
  Frac x == Frac y = x == y
  Inter x == Frac y = fromIntegral x == y
  x == y = y == x

instance Ord Number where
  Inter x <= Inter y = x <= y
  Frac x <= Frac y = x <= y
  Inter x <= Frac y = fromIntegral x <= y
  Frac x <= Inter y = x <= fromIntegral y

data Val
  = Boolean Bool
  | Str String
  | Numb Number
  deriving (Show)

add :: Val -> Val -> Val
add (Boolean a) (Boolean b) = Boolean $ a || b
add (Str a) (Str b)         = Str $ a <> b
add (Numb a) (Numb b)       = Numb $ a <> b
add _ _                     = Str mempty

greater :: Val -> Val -> Bool
greater (Boolean a) (Boolean b) = a > b
greater (Str a) (Str b)         = a > b
greater (Numb a) (Numb b)       = a > b
greater _ _                     = False

class ToVal a where
  toVal :: a -> Val

instance ToVal Bool where
  toVal = Boolean

instance ToVal Int where
  toVal = Numb . Inter

instance ToVal Double where
  toVal = Numb . Frac

instance ToVal String where
  toVal = Str

class Monad m =>
      MonadJS m s
  | m -> s
  where
  type Var m a :: *
  newVar :: a -> m (Var m a)
  writeVar :: Var m a -> a -> m ()
  (@>@) :: m Val -> m Val -> m Bool
  (@#) :: m () -> m a -> m a
  (@+@) :: m Val -> m Val -> m Val
  (@=@) :: m (Var m Val) -> m Val -> m ()
  sReadVar :: m (Var m Val) -> m Val
  sIf :: m Bool -> m () -> m () -> m ()
  sFun1 :: (m Val -> m (Var m Val) -> m ()) -> m Val -> m Val
  sWhile :: m Bool -> m () -> m ()

class VarJS a where
  sWithVar :: MonadJS m s => a -> (m (Var m Val) -> m ()) -> m ()
  (@=) :: MonadJS m s => m (Var m Val) -> a -> m ()

instance MonadJS (ST s) s where
  type Var (ST s) a = STRef s a
  newVar = newSTRef
  writeVar = writeSTRef
  (@>@) = liftA2 greater
  (@#) = (>>)
  (@+@) = liftA2 add
  (@=@) mRef mX = do
    x <- mX
    ref <- mRef
    writeSTRef ref x
  sReadVar mVar = do
    var <- mVar
    readSTRef var
  sIf cond a b = do
    flag <- cond
    if flag
      then a
      else b
  sFun1 cont arg = do
    resRef <- newSTRef $ Str ""
    cont arg $ return resRef
    readSTRef resRef
  sWhile cond body = do
    flag <- cond
    when flag $ body >> sWhile cond body

instance VarJS Val where
  sWithVar x cont = do
    ref <- newVar x
    cont (return ref)
  (@=) mRef x = do
    ref <- mRef
    writeVar ref x

sWithVarGeneric ::
     (MonadJS m s, ToVal a) => a -> (m (Var m Val) -> m ()) -> m ()
sWithVarGeneric x = sWithVar $ toVal x

assignGeneric :: (MonadJS m s, ToVal a) => m (Var m Val) -> a -> m ()
assignGeneric mRef x = mRef @= toVal x

instance VarJS Int where
  sWithVar = sWithVarGeneric
  (@=) = assignGeneric

instance VarJS Double where
  sWithVar = sWithVarGeneric
  (@=) = assignGeneric

instance VarJS Bool where
  sWithVar = sWithVarGeneric
  (@=) = assignGeneric

instance VarJS String where
  sWithVar = sWithVarGeneric
  (@=) = assignGeneric

runSTJS :: (forall s. ST s ()) -> ()
runSTJS = runST

runSTJS1 :: ToVal a => a -> (forall s. ST s Val -> ST s Val) -> Val
runSTJS1 x script = runST $ script (return $ toVal x)

f1 :: MonadJS m s => m Val -> m Val
f1 = sFun1 $ \a res -> res @=@ a

f2 :: MonadJS m s => m ()
f2 = sWithVar (toVal "keke") $ \x -> x @= toVal False

f3 :: MonadJS m s => m Val -> m Val
f3 =
  sFun1 $ \x res ->
    sWithVar (2 :: Int) $ \y ->
      sIf (sReadVar y @>@ x) (res @=@ x) (res @= toVal "Ochen zhal")
-- f :: ToVal a => a -> exists b . VarJS b
-- instance ToVal a => VarJS a where
--   sWithVar x cont = sWithVar (toVal x) cont
--   (@=) mRef x = mRef @= toVal x
-- data JS s r where
--   PureVal :: VarCreator a => a -> JS s Var
--   JSVar :: STRef s Var -> JS s (STRef s Var)
--   WithVar :: VarBox -> (JS s (STRef s Var) -> JS s ()) -> JS s ()
--   (:=) :: JS s (STRef s Var) -> VarBox -> JS s ()
--   (:#) :: JS s () -> JS s a -> JS s a
--   (:+) :: JS s Var -> JS s Var -> JS s Var
--   Fun1 :: (JS s Var -> JS s (STRef s Var) -> JS s ()) -> JS s Var -> JS s Var
--   ReadVar :: JS s (STRef s Var) -> JS s Var
--   (:>) :: JS s Var -> JS s Var -> JS s Bool
--   If :: JS s Bool -> JS s () -> JS s () -> JS s ()
--   (:<=>) :: JS s (STRef s Var) -> Var -> JS s ()
-- jsToST :: forall s r . JS s r -> ST s r
-- jsToST (JSVar x) = pure x
-- jsToST (WithVar (VB initVal) comp) = do
--   var <- newSTRef $ mkVar initVal
--   jsToST $ comp $ JSVar var
-- jsToST (t := (VB x)) = do
--   ref <- jsToST t
--   writeSTRef ref $ mkVar x
-- jsToST (a :# b) = jsToST a >> jsToST b
-- jsToST (a :+ b) = do
--   x <- jsToST a
--   y <- jsToST b
--   return $ add x y
-- jsToST (Fun1 comp arg) = do
--   resRef <- newSTRef $ Str ""
--   let jsFunc = comp arg $ JSVar resRef
--   jsToST jsFunc
--   readSTRef resRef
-- jsToST (ReadVar var) = do
--   ref <- jsToST var
--   readSTRef ref
-- jsToST (a :> b) = do
--   x <- jsToST a
--   y <- jsToST b
--   return $ greater x y
-- jsToST (If cond a b) = do
--   flag <- jsToST cond
--   if flag
--     then jsToST a
--     else jsToST b
-- jsToST (PureVal x) = return $ mkVar x
-- jsToST (t :<=> x) = do
--   ref <- jsToST t
--   writeSTRef ref x
-- run :: (forall s . JS s ()) -> ()
-- run script = runST $ jsToST script
-- run1 :: VarCreator a => a -> (forall s . JS s Var -> JS s Var) -> Var
-- run1 a script = runST $ jsToST $ script (PureVal a)
-- f1 = Fun1 $ \a res -> res := (VB (2 :: Int))