{-# LANGUAGE PolyKinds, TypeFamilies, AllowAmbiguousTypes, FlexibleContexts, DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses, TypeOperators, ConstraintKinds, ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications, GADTs, UndecidableInstances, UndecidableSuperClasses #-}
{-# LANGUAGE RankNTypes, DefaultSignatures, FlexibleInstances, EmptyCase, InstanceSigs #-}
module Data.Serialize.Versioned.Internal ( VersionDomain(..)
                                         , VersionedSerialize(..)
                                         , VersionedGettable(..)
                                         , MigrationType(..)
                                         , MigrationTypeWitness(..)
                                         , IsMigrationType
                                         , migrationTypeWitness
                                         , MigrateType
                                         , ConsistentMigrationVariant
                                         , UnchangedRepresentation
                                         , PreviousRepresentation
                                         , IsPreviousRepresentation
                                         , HasPreviousVersion
                                         , VersionedGetType
                                         , ValidVersion
                                         , DefinedVersion(..)
                                         , ForcePreviousVersion
                                         , FromJust
                                         , DefaultVersionedGet(..)
                                         , GVersionedGet(..)
                                         , GVersionedPut(..)
                                         , versionedGetWithVersion
                                         , getVersionWord
                                         , getUnversionedResumable
                                         , getUnversioned
                                         , getVersioned
                                         , label
                                         , putUnversioned
                                         ) where

import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Reader ( ReaderT(ReaderT)
                                  , runReaderT
                                  , ask
                                  )
import Data.Kind ( Constraint
                 , Type
                 )
import Data.Proxy (Proxy(Proxy))
import qualified Data.Serialize as Cereal
import Data.Serialize ( Get
                      , PutM
                      , Serialize
                      , get
                      , put
                      )
import Data.Word ( Word8
                 , Word16
                 , Word32
                 , Word64
                 )
import GHC.Generics ( (:+:)(L1, R1)
                    , (:*:)((:*:))
                    , C1
                    , Constructor
                    , D1
                    , Datatype
                    , Generic
                    , K1(K1)
                    , M1(M1)
                    , Rep
                    , S1
                    , Selector
                    , U1(U1)
                    , V1
                    , conName
                    , datatypeName
                    , from
                    , selName
                    , to
                    )
import GHC.TypeLits ( type (^)
                    , type (-)
                    , CmpNat
                    , KnownNat
                    , Nat
                    , TypeError
                    , ErrorMessage ( (:<>:)
                                   , ShowType
                                   , Text
                                   )
                    , natVal
                    )

import Data.Serialize.Versioned.Get.Types
import Data.Serialize.Versioned.Put.Types

-- | Marks a type (of any kind) as being usable as a domain.
class DefinedVersion d (CurrentVersion d) => VersionDomain (d :: dk) where
  -- | The current version in the domain.
  type CurrentVersion d :: Nat

-- | Versions are serialized to disk as 64 bit numbers. This
-- constraint ensures that the version fits that size. It also ensures
-- that 2^64-1 is not used as a version, because it is reserved for
-- future use.
type ValidVersion v = (CmpNat v ((2 ^ 64) - 1) ~ 'LT, KnownNat v)

-- | Defines the preceeding version of @v@ (possibly @'Nothing@) in
-- the versioning domain @d@.
class ValidVersion v => DefinedVersion d v where
  type PreviousVersion d v :: Maybe Nat

-- | This class encodes knowledge on how to deserialize a data type
-- @t@ in a given domain @d@ and version @v@ and how to migrate it
-- from a previous version if necessary.
class ( DefinedVersion d v
      , IsMigrationType (MigrationVariant d v t)
      , ConsistentMigrationVariant (MigrationVariant d v t) d v t
      ) => VersionedGettable (d :: dk) (v :: Nat) (t :: *) where
  -- | Defines the relation between the structure in this version and
  -- its predecessor in the previous version.
  type MigrationVariant d v t :: MigrationType
  -- | Defines how to migrate a value of the previous version to this
  -- version.
  --
  --The type depends on the 'MigrationVariant'.
  --
  -- A default implementation is available if the 'MigrationVariant'
  -- is 'Base'.
  migrate :: MigrateType (MigrationVariant d v t) t
  default migrate :: (MigrateType (MigrationVariant d v t) t ~ ())
                  => MigrateType (MigrationVariant d v t) t
  migrate = ()
  -- | Defines how to deserialize the type @t@ without worriyng about
  -- version tags or migration. You should only define this in your
  -- own instances, not use it! Use
  -- 'Data.Serialize.Versioned.Get.getVersioned' instead.
  --
  -- For 'Generic' types a default implementation can always be
  -- generated. Otherwise a default implementation is only available
  -- when 'MigrationVariant' is 'Unchanged'.
  versionedGet :: VersionedGetType (MigrationVariant d v t) d t
  default versionedGet :: (DefaultVersionedGet (MigrationVariant d v t) d v t)
                       => VersionedGetType (MigrationVariant d v t) d t
  versionedGet = defaultVersionedGet @_ @d @v @t

-- | This class marks the current version of a structure and how to
-- serialize it for the current version. Note that even though
-- deserializing old data is possible, only current structures can be
-- encoded to data of current versions.
class ( VersionedGettable d (CurrentStructureVersion d t) t
      , VersionDomain d
      ) => VersionedSerialize (d :: dk) (t :: *) where
  -- | Defines the current version of @t@ in domain @d@. Defaults to
  -- @'CurrentVersion' d@.
  type CurrentStructureVersion d t :: Nat
  type CurrentStructureVersion d t = CurrentVersion d
  -- | Defines how to serialize the type @t@ without worriyng about
  -- version tags or migration.
  --
  -- A default implementation is available for 'Generic' types.
  putVersioned :: VersionedPutter d (CurrentStructureVersion d t) t
  default putVersioned :: ( GVersionedPut d (CurrentStructureVersion d t) (Rep t)
                          , Generic t
                          )
                       => VersionedPutter d (CurrentStructureVersion d t) t
  putVersioned = gVersionedPut . from

-- | This type is intended to be used as a DataKind. It specifies the
-- relation between a data type with a given version and its
-- predecessor in the previous version.
data MigrationType = Base
                   -- ^ The structure had has no predecessor in the
                   -- previous version, possibly because there is no
                   -- previous version.
                   | Unchanged
                   -- ^ The structure was represented and serialized
                   -- the same way in the previous version.
                   | SerializationChanged
                   -- ^ The structure was represented the same way
                   -- (its datatype didn't change) in the previous
                   -- version, but the way it is serialized changed.
                   | RepresentationChanged Type
                   -- ^ The structure is represented differently than
                   -- in the previous version

data MigrationTypeWitness (mt :: MigrationType) where
  BaseWitness :: MigrationTypeWitness 'Base
  UnchangedWitness :: MigrationTypeWitness 'Unchanged
  SerializationChangedWitness :: MigrationTypeWitness 'SerializationChanged
  RepresentationChangedWitness :: MigrationTypeWitness ('RepresentationChanged t)

-- | This class is used internally to work with
-- 'MigrationType's. Notably it is a superclass of
-- 'VersionedSerialize'. As it is implemented for all types of kind
-- 'MigrationType' it should always be satisfied.
class IsMigrationType (mt :: MigrationType) where
  migrationTypeWitness :: MigrationTypeWitness mt

instance IsMigrationType 'Base where
  migrationTypeWitness = BaseWitness

instance IsMigrationType 'Unchanged where
  migrationTypeWitness = UnchangedWitness

instance IsMigrationType 'SerializationChanged where
  migrationTypeWitness = SerializationChangedWitness

instance IsMigrationType ('RepresentationChanged t) where
  migrationTypeWitness = RepresentationChangedWitness

-- | This type family imposes superclasses for 'VersionedSerialize'
-- depending on the defined 'MigrationVariant' to ensure, that the
-- claim you make by defining the 'MigrationVariant' is actually
-- consistent with other observations.
type family ConsistentMigrationVariant (mt :: MigrationType) (d :: dk) (v :: Nat) (t :: *) :: Constraint where
  ConsistentMigrationVariant 'Base _ _ _ = ()
  ConsistentMigrationVariant 'Unchanged d v t = UnchangedRepresentation d v t
  ConsistentMigrationVariant 'SerializationChanged d v t = UnchangedRepresentation d v t
  ConsistentMigrationVariant ('RepresentationChanged oldT) d v newT =
    IsPreviousRepresentation d v newT oldT

-- | This constraint ensures that the previous representation of
-- @newT@ is @oldT@ in the previous version.
type IsPreviousRepresentation d v newT oldT =
  ( HasPreviousVersion d v
  , VersionedGettable d (ForcePreviousVersion d v) oldT
  , MigrationVariant d v newT ~ 'RepresentationChanged oldT
  )

type PreviousRepresentation d v t = PreviousRepresentationGo (MigrationVariant d v t) d v t

type family PreviousRepresentationGo (mt :: MigrationType) (d :: dk) (v :: Nat) (t :: *) where
  PreviousRepresentationGo ('RepresentationChanged t) _ _ _ = t
  PreviousRepresentationGo _ d v t =
    TypeError ('Text "PreviousRepresentation: " ':<>: 'ShowType t ':<>: 'Text " has no previous representation relative to " ':<>: 'ShowType v ':<>: 'Text " in domain " ':<>: 'ShowType d)

-- | This constraint that ensures that the representation for @t@ hasn't
-- changed since the previous version.
type UnchangedRepresentation d v t =
  ( VersionedGettable d (ForcePreviousVersion d v) t
  , HasPreviousVersion d v
  )

-- | This constraint ensures that @v@ has a preceeding version in
-- domain @d@
type HasPreviousVersion (d :: dk) (v :: Nat) =
  PreviousVersion d v ~ 'Just (FromJust (PreviousVersion d v))

-- | The previous version of @v@ in the domain @d@ or a 'TypeError' if
-- there is no such preceeding version. Should normally only be used
-- when @'HasPreviousVersion' d v@ holds.
type ForcePreviousVersion (d :: dk) (v :: Nat) = FromJust (PreviousVersion d v)

-- | 'Data.Maybe.fromJust' on type level
type family FromJust (m :: Maybe dk) :: dk where
  FromJust 'Nothing = TypeError ('Text "FromJust: Nothing")
  FromJust ('Just t) = t

-- | This defines the type for 'migrate' for each possible definition
-- of 'MigrationVariant'.
--
-- The type of 'migrate' is @oldT -> newT@ if @'MigrationVariant' tag
-- newT@ is @'RepresentationChanged' oldT@, otherwise it is @()@.
type family MigrateType (mt :: MigrationType) (t :: *) where
  MigrateType ('RepresentationChanged oldT) newT =
    oldT -> newT
  MigrateType _ _ = ()

-- | This defines the type for 'versionedGet' for each possible
-- definition of 'MigrationVariant'.
--
-- The type of 'versionedGet' is @()@ if @'MigrationVariant' tag t@ is
-- 'Unchanged', otherwise it is @'VersionedGet' tag t@
type family VersionedGetType (mt :: MigrationType) (d :: dk) (t :: *) where
  VersionedGetType 'Unchanged _ _ = ()
  VersionedGetType _ d t = VersionedGet d t

getVersionWord :: forall (v :: Nat). (ValidVersion v) => Word64
getVersionWord = fromIntegral $ natVal (Proxy @v)

-- | This class provides a default implementation for
-- 'versionedGet'.
class (mt ~ MigrationVariant d v t) => DefaultVersionedGet mt d v t where
  -- | A default implementation for 'versionedGet'. Either @()@ or
  -- 'gVersionedGet', depending on the required type.
  defaultVersionedGet :: VersionedGetType (MigrationVariant d v t) d t

instance ( MigrationVariant d v t ~ 'Base
         , GVersionedGet d (Rep t)
         , Generic t
         ) => DefaultVersionedGet 'Base d v t where
  defaultVersionedGet = fmap to gVersionedGet

instance ( MigrationVariant d v t ~ 'Unchanged
         ) => DefaultVersionedGet 'Unchanged d v t where
  defaultVersionedGet = ()

instance ( MigrationVariant d v t ~ 'SerializationChanged
         , GVersionedGet d (Rep t)
         , Generic t
         ) => DefaultVersionedGet 'SerializationChanged d v t where
  defaultVersionedGet = fmap to gVersionedGet

instance ( MigrationVariant d v t ~ 'RepresentationChanged oldT
         , GVersionedGet d (Rep t)
         , Generic t
         ) => DefaultVersionedGet ('RepresentationChanged oldT) d v t where
  defaultVersionedGet = fmap to gVersionedGet

versionedGetWithVersion :: VersionedGet d a -> Word64 -> Get a
versionedGetWithVersion (VersionedGet vg) = runReaderT vg

-- | A version of getUnversioned that allows you to resume decoding
-- versioned data inside the unversioned data, without parsing the
-- version tag again.
getUnversionedResumable :: ((forall b. VersionedGet d b -> Get b) -> Get a)
                        -> VersionedGet d a
getUnversionedResumable f =
  VersionedGet $ ReaderT $ \v -> f $ flip versionedGetWithVersion v

-- | Labels a block of code for debugging. Direct lift of
-- 'Cereal.label'.
label :: String -> VersionedGet d a -> VersionedGet d a
label l x = getUnversionedResumable $ \resume -> Cereal.label l $ resume x

-- | Lift an unversioned 'Get' operation into 'VersionedGet'. The
-- lifted operation will not be able to know the on-disk version or
-- the target version. You can however lift another call to
-- 'Data.Serialize.Versioned.Get.runVersionedGet' to deserialize a
-- seperately versioned container that is embedded in the outer
-- container.
getUnversioned :: Get a -> VersionedGet d a
getUnversioned = VersionedGet . lift

-- | Reads a value of the required type or one if its predecessor
-- types and migrates it to the target version and type.
getVersioned :: forall d a.
                (VersionedSerialize d a)
             => VersionedGet d a
getVersioned = getVersioned' @_ @(CurrentStructureVersion d a)

getVersioned' :: forall d v a.
                 (VersionedGettable d v a)
              => VersionedGet d a
getVersioned' = go @_ @v False
  where go :: forall d' (v' :: Nat) a'.
              (VersionedGettable d' v' a')
           => Bool -> VersionedGet d' a'
        go alreadyHasCurrentVersion = do
          hasCurrentVersion <- (|| alreadyHasCurrentVersion) . (==getVersionWord @v') <$>
                               VersionedGet ask
          case migrationTypeWitness @(MigrationVariant d' v' a') of
            BaseWitness ->
              if hasCurrentVersion
              then getSelf
              else fail migrationBaseErrorMessage
            UnchangedWitness -> getPreviousSame hasCurrentVersion
            SerializationChangedWitness ->
              if hasCurrentVersion
              then getSelf
              else getPreviousSame hasCurrentVersion
            RepresentationChangedWitness ->
              if hasCurrentVersion
              then getSelf
              else migrate @_ @d' @v' @a' <$>
                   go @_ @(ForcePreviousVersion d' v') hasCurrentVersion
            where 
              migrationBaseErrorMessage =
                "getVersioned: Unmigratable version encountered"
              getSelf = versionedGet @_ @d' @v' @a'
              getPreviousSame :: (UnchangedRepresentation d' v' a')
                              => Bool -> VersionedGet d' a'
              getPreviousSame =
                go @_ @(ForcePreviousVersion d' v')

-- | Lifts a 'PutM' operation into 'VersionedPutM'. In other word,
-- puts a bit of unversioned information into a versioned container.
--
-- Note that you can call
-- 'Data.Serialize.Versioned.Put.runVersionedPut' in the inner block
-- again, to embed a separately versioned container into another.
putUnversioned :: PutM a -> VersionedPutM d v a
putUnversioned = VersionedPutM

-- | Provides a default implementation for deserializing generic
-- structures.
class GVersionedGet (d :: dk) f where
  -- | Deserializes a generic structure.
  gVersionedGet :: VersionedGet d (f a)

instance GVersionedGet d V1 where
  gVersionedGet = fail "gVersionedGet: V1"

instance GVersionedGet d U1 where
  gVersionedGet = pure U1

instance (VersionedSerialize d a) => GVersionedGet d (K1 i a) where
  gVersionedGet = fmap K1 getVersioned

instance ( GVersionedGet d f
         , Datatype t
         ) => GVersionedGet d (D1 t f) where
  gVersionedGet :: forall a. VersionedGet d (D1 t f a)
  gVersionedGet =
    labeledGVersionedGet $ "gVersionedGet for type " ++ datatypeName (undefined :: D1 t f a)

instance ( GVersionedGet d f
         , Constructor c
         ) => GVersionedGet d (C1 c f) where
  gVersionedGet :: forall a. VersionedGet d (C1 c f a)
  gVersionedGet =
    labeledGVersionedGet $ "gVersionedGet for constructor " ++ conName (undefined :: C1 c f a)

instance ( GVersionedGet d f
         , Selector s
         ) => GVersionedGet d (S1 s f) where
  gVersionedGet :: forall a. VersionedGet d (S1 s f a)
  gVersionedGet =
      let sn = selName (undefined :: S1 s f a)
          l = if null sn
              then "unlabeled constructor field"
              else "selector field " ++ sn
      in labeledGVersionedGet $ "gVersionedGet for " ++ l

labeledGVersionedGet :: (GVersionedGet d f)
                      => String
                      -> VersionedGet d (M1 i c f a)
labeledGVersionedGet s = M1 <$> label s gVersionedGet

instance ( GVersionedGet d f
         , GVersionedGet d g
         ) => GVersionedGet d (f :*: g) where
  gVersionedGet = (:*:) <$> gVersionedGet <*> gVersionedGet

instance ( GVersionedGetSum d f
         , GVersionedGetSum d g
         , SumSize f
         , SumSize g
         ) => GVersionedGet d (f :+: g) where
  gVersionedGet
    | size <= maxBoundI @Word8 = getSumStart @Word8
    | size <= maxBoundI @Word16 = getSumStart @Word16
    | size <= maxBoundI @Word32 = getSumStart @Word32
    | size <= maxBoundI @Word64 = getSumStart @Word64
    | otherwise = sizeError "decode" size
    where size = sumSize @(f :+: g)

-- | Provides a default implementation for serializing generic
-- structures.
class GVersionedPut (d :: dk) (v :: Nat) f where
  -- | Serializes a generic structure.
  gVersionedPut :: VersionedPutter d v (f a)

instance GVersionedPut d v V1 where
  gVersionedPut v = case v of {}

instance GVersionedPut d v U1 where
  gVersionedPut U1 = pure ()

instance ( VersionedSerialize d a
         , CurrentStructureVersion d a ~ v
         ) => GVersionedPut d v (K1 i a) where
  gVersionedPut (K1 x) = putVersioned x

instance (GVersionedPut d v f) => GVersionedPut d v (M1 i c f) where
  gVersionedPut (M1 x) = gVersionedPut x

instance ( GVersionedPut d v f
         , GVersionedPut d v g
         ) => GVersionedPut d  v(f :*: g) where
  gVersionedPut (x :*: y) = gVersionedPut x <> gVersionedPut y

instance ( GVersionedPutSum d v f
         , GVersionedPutSum d v g
         , SumSize f
         , SumSize g
         ) => GVersionedPut d v (f :+: g) where
  gVersionedPut
    | size <= maxBoundI @Word8 = putSum (0 :: Word8)
    | size <= maxBoundI @Word16 = putSum (0 :: Word16)
    | size <= maxBoundI @Word32 = putSum (0 :: Word32)
    | size <= maxBoundI @Word64 = putSum (0 :: Word64)
    | otherwise = sizeError "encode" size
    where size = sumSize @(f :+: g)

sizeError :: (Show s) => String -> s -> a
sizeError n s = error $ "Can't " ++ n ++ " type with " ++ show s ++ " constructors."

getSumStart :: forall word d f a.
               ( GVersionedGetSum d f
               , SumSize f
               , Integral word
               , Serialize word
               )
            => VersionedGet d (f a)
getSumStart = do
  code <- getUnversioned (get @word)
  if fromIntegral code > sumSize @f
    then fail "Unknown constructor encountered in encoded data"
    else getSum code

maxBoundI :: forall b. (Bounded b, Integral b) => Integer
maxBoundI = fromIntegral $ maxBound @b

class GVersionedGetSum d f where
  getSum :: (Integral word, Serialize word) => word -> VersionedGet d (f a)

instance ( GVersionedGetSum d f
         , GVersionedGetSum d g
         , SumSize f
         ) => GVersionedGetSum d (f :+: g) where
  getSum code
    | fromIntegral code < sumSize @f = L1 <$> getSum code
    | otherwise = R1 <$> getSum code

instance ( GVersionedGet d a
         , Constructor c
         ) => GVersionedGetSum d (C1 c a) where
  getSum _ = gVersionedGet

class GVersionedPutSum d v f where
  putSum :: (Integral word, Serialize word)
         => word -> VersionedPutter d v (f a)

instance ( GVersionedPutSum d v f
         , GVersionedPutSum d v g
         , SumSize f
         ) => GVersionedPutSum d v (f :+: g) where
  putSum offset (L1 x) = putSum offset x
  putSum offset (R1 x) = putSum (offset + fromInteger (sumSize @f)) x

instance (GVersionedPut d v a) => GVersionedPutSum d v (C1 c a) where
  putSum code x = putUnversioned (put code) *> gVersionedPut x

class SumSize f where
  sumSize :: Integer

instance (SumSize f, SumSize g) => SumSize (f :+: g) where
  sumSize = sumSize @f + sumSize @g

instance SumSize (C1 c a) where
  sumSize = 1