-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | General type utilities.
module Morley.Util.Type
  ( type (==)
  , If
  , type (++)
  , IsElem
  , type (/)
  , type (//)
  , Guard
  , FailWhen
  , FailUnless
  , failUnlessEvi
  , failWhenEvi
  , AllUnique
  , RequireAllUnique
  , ReifyList (..)
  , PatternMatch
  , PatternMatchL
  , KnownList (..)
  , KList (..)
  , RSplit
  , rsplit
  , Some1 (..)
  , recordToSomeList

  , reifyTypeEquality

  , ConcatListOfTypesAssociativity
  , listOfTypesConcatAssociativityAxiom

  , MockableConstraint (..)

  , onFirst
  , knownListFromSingI
  ) where

import Data.Constraint (Dict(..), (:-)(..), (\\))
import Data.Singletons (SingI(sing))
import Data.Singletons.Prelude.List (SList(..))
import Data.Type.Bool (If, Not, type (&&))
import Data.Type.Equality (type (==))
import Data.Vinyl.Core (Rec(..))
import Data.Vinyl.Functor qualified as Vinyl
import Data.Vinyl.Recursive (recordToList, rmap)
import Data.Vinyl.TypeLevel (type (++))
import GHC.TypeLits (ErrorMessage(..), Symbol, TypeError)
import Unsafe.Coerce (unsafeCoerce)

type family IsElem (a :: k) (l :: [k]) :: Bool where
  IsElem _ '[] = 'False
  IsElem a (a ': _) = 'True
  IsElem a (_ ': as) = IsElem a as

-- | Remove all occurences of the given element.
type family (l :: [k]) / (a :: k) where
  '[] / _ = '[]
  (a ': xs) / a = xs / a
  (b ': xs) / a = b ': (xs / a)

-- | Difference between two lists.
type family (l1 :: [k]) // (l2 :: [k]) :: [k] where
  l // '[] = l
  l // (x ': xs) = (l / x) // xs

type family Guard (cond :: Bool) (a :: k) :: Maybe k where
  Guard 'False _ = 'Nothing
  Guard 'True a = 'Just a

-- | Fail with given error if the condition does not hold.
type family FailUnless (cond :: Bool) (msg :: ErrorMessage) :: Constraint where
  FailUnless 'True _ = ()
  FailUnless 'False msg = TypeError msg

-- | Fail with given error if the condition holds.
type FailWhen cond msg = FailUnless (Not cond) msg

-- | A natural conclusion from the fact that an error has not occurred.
failUnlessEvi :: forall cond msg. FailUnless cond msg :- (cond ~ 'True)
failUnlessEvi = Sub unsafeProvideConstraint

failWhenEvi :: forall cond msg. FailWhen cond msg :- (cond ~ 'False)
failWhenEvi = Sub unsafeProvideConstraint

type family AllUnique (l :: [k]) :: Bool where
  AllUnique '[] = 'True
  AllUnique (x : xs) = Not (IsElem x xs) && AllUnique xs

type RequireAllUnique desc l = RequireAllUnique' desc l l

type family RequireAllUnique' (desc :: Symbol) (l :: [k]) (origL ::[k]) :: Constraint where
  RequireAllUnique' _ '[] _ = ()
  RequireAllUnique' desc (x : xs) origL =
    If (IsElem x xs)
       (TypeError ('Text "Duplicated " ':<>: 'Text desc ':<>: 'Text ":" ':$$:
                   'ShowType x ':$$:
                   'Text "Full list: " ':<>:
                   'ShowType origL
                  )
       )
       (RequireAllUnique' desc xs origL)

-- | Make sure given type is evaluated.
-- This type family fits only for types of 'Type' kind.
type family PatternMatch (a :: Type) :: Constraint where
  PatternMatch Int = ((), ())
  PatternMatch _ = ()

type family PatternMatchL (l :: [k]) :: Constraint where
  PatternMatchL '[] = ((), ())
  PatternMatchL _ = ()

-- | Bring type-level list at term-level using given function
-- to demote its individual elements.
class ReifyList (c :: k -> Constraint) (l :: [k]) where
  reifyList :: (forall a. c a => Proxy a -> r) -> [r]

instance ReifyList c '[] where
  reifyList _ = []

instance (c x, ReifyList c xs) => ReifyList c (x ': xs) where
  reifyList reifyElem = reifyElem (Proxy @x) : reifyList @_ @c @xs reifyElem

-- | Reify type equality from boolean equality.
reifyTypeEquality :: forall a b x. (a == b) ~ 'True => (a ~ b => x) -> x
reifyTypeEquality x = x \\ unsafeProvideConstraint @(a ~ b)

-- | Similar to @SingI []@, but does not require individual elements to be also
-- instance of @SingI@.
class KnownList l where
  klist :: KList l
instance KnownList '[] where
  klist = KNil
instance KnownList xs => KnownList (x ': xs) where
  klist = KCons Proxy Proxy

-- | Given a type-level list that is 'SingI', construct evidence that it is also
-- a 'KnownList'. Note that 'KnownList' is weaker, hence this construction
-- is always possible.
knownListFromSingI :: forall xs. SingI xs => Dict (KnownList xs)
knownListFromSingI = go (sing @xs)
  where
    go :: forall ys. SList ys -> Dict (KnownList ys)
    go = \case
      SNil -> Dict
      SCons _ ys -> Dict \\ go ys

-- | @Data.List.Singletons.SList@ analogue for 'KnownList'.
data KList (l :: [k]) where
  KNil :: KList '[]
  KCons :: KnownList xs => Proxy x -> Proxy xs -> KList (x ': xs)

type RSplit l r = KnownList l

-- | Split a record into two pieces.
rsplit
  :: forall k (l :: [k]) (r :: [k]) f.
      (RSplit l r)
  => Rec f (l ++ r) -> (Rec f l, Rec f r)
rsplit = case klist @l of
  KNil -> (RNil, )
  KCons{} -> \(x :& r) ->
    let (x1, r1) = rsplit r
    in (x :& x1, r1)

-- | A value of type parametrized with /some/ type parameter.
data Some1 (f :: k -> Type) =
  forall a. Some1 (f a)

deriving stock instance (forall a. Show (f a)) => Show (Some1 f)

recordToSomeList :: Rec f l -> [Some1 f]
recordToSomeList = recordToList . rmap (Vinyl.Const . Some1)

type ConcatListOfTypesAssociativity a b c = ((a ++ b) ++ c) ~ (a ++ (b ++ c))

-- | GHC can't deduce this itself because
-- in general a type family might be not associative,
-- what brings extra difficulties and redundant constraints,
-- especially if you have complex types.
-- But (++) type family is associative, so let's define this small hack.
listOfTypesConcatAssociativityAxiom :: forall a b c . Dict (ConcatListOfTypesAssociativity a b c)
listOfTypesConcatAssociativityAxiom = unsafeProvideConstraint

-- | Constaints that can be provided on demand.
--
-- Needless to say, this is a pretty unsafe operation. This typeclass makes
-- using it safer in a sense that getting a segfault becomes harder, but still
-- it deceives the type system and should be used only if providing a proper
-- proof would be too difficult.
class MockableConstraint (c :: Constraint) where
  -- | Produce a constraint out of thin air.
  unsafeProvideConstraint :: Dict c

instance MockableConstraint (a ~ b) where
  unsafeProvideConstraint = unsafeCoerce $ Dict @(Int ~ Int)

-- | In majority of the cases we don't want to mock typeclass constraints
-- since coercing instances of typeclasses with methods is utterly unsafe.
instance {-# OVERLAPPABLE #-} c a => MockableConstraint (c a) where
  unsafeProvideConstraint = Dict

instance ( MockableConstraint c1
         , MockableConstraint c2
         ) => MockableConstraint (c1, c2) where
  unsafeProvideConstraint = runIdentity $ do
    Dict <- pure $ unsafeProvideConstraint @c1
    Dict <- pure $ unsafeProvideConstraint @c2
    return Dict
  {-# INLINE unsafeProvideConstraint #-}

instance ( MockableConstraint c1
         , MockableConstraint c2
         , MockableConstraint c3
         ) => MockableConstraint (c1, c2, c3) where
  unsafeProvideConstraint = runIdentity $ do
    Dict <- pure $ unsafeProvideConstraint @(c1, c2)
    Dict <- pure $ unsafeProvideConstraint @c3
    return Dict
  {-# INLINE unsafeProvideConstraint #-}

instance ( MockableConstraint c1
         , MockableConstraint c2
         , MockableConstraint c3
         , MockableConstraint c4
         ) => MockableConstraint (c1, c2, c3, c4) where
  unsafeProvideConstraint = runIdentity $ do
    Dict <- pure $ unsafeProvideConstraint @(c1, c2, c3)
    Dict <- pure $ unsafeProvideConstraint @c4
    return Dict
  {-# INLINE unsafeProvideConstraint #-}

instance ( MockableConstraint c1
         , MockableConstraint c2
         , MockableConstraint c3
         , MockableConstraint c4
         , MockableConstraint c5
         ) => MockableConstraint (c1, c2, c3, c4, c5) where
  unsafeProvideConstraint = runIdentity $ do
    Dict <- pure $ unsafeProvideConstraint @(c1, c2, c3, c4)
    Dict <- pure $ unsafeProvideConstraint @c5
    return Dict
  {-# INLINE unsafeProvideConstraint #-}

-- | Utility function to help transform the first argument of a binfunctor.
onFirst :: Bifunctor p => p a c -> (a -> b) -> p b c
onFirst = flip first