{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
module Pretty.Serialize (Serializer, PrettySerialize(..), ShowPrettySerialize(..), run, emit, KeyValueSerialize(..), customField) where

import Data.Text.Lazy (LazyText)
import GHC.Generics (U1, K1 (K1), Generically (Generically), Generic (Rep), from, M1 (M1), Meta(MetaSel, MetaCons, MetaData), (:+:) (L1, R1), (:*:) ((:*:)))
import Control.Monad ((<$!>), forM_)
import qualified Data.Text as StrictText
import Data.String (IsString(fromString))
import GHC.TypeLits (symbolVal, KnownSymbol)
import Data.Proxy (Proxy(Proxy))
import Data.Word (Word8, Word16, Word32, Word64)
import Data.Int (Int8, Int16, Int32, Int64)
import Numeric.Natural (Natural)
import Data.Functor.Identity (Identity)
import Data.Functor.Const (Const)
import Data.Functor.Compose (Compose)
import Data.Kind (Type)
import Data.List (List)
import Data.Typeable (typeOf, Typeable)
import qualified Data.Array.IArray as Data.IArray
import Data.Tuple (Solo)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.ByteString.Lazy (LazyByteString)
import Data.ByteString (ByteString)
import Data.Maybe (fromMaybe)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Ratio (Ratio)
import Data.Text (StrictText)
import Data.Bifunctor (Bifunctor(second))
import qualified Data.Text.Lazy as LazyText

type Indentation = Word

newtype Serializer a = Serializer (ContinuationSerializer a)
  deriving stock (Functor)

type ContinuationSerializer a = (SerializerState -> (SerializerState, a))

type Prepends = [StrictText] -> [StrictText]

data SerializerState = SerializerState
  { builder             :: Prepends
  , indentation         :: !Indentation
  , increaseIndentation :: !(Indentation -> Indentation)
  , currentFieldIndex   :: !(Maybe Word)
  }

instance Applicative Serializer where
  pure :: a -> Serializer a
  pure x = Serializer (, x)
  (<*>) :: Serializer (a -> b) -> Serializer a -> Serializer b
  (<*>) (Serializer computeF) (Serializer computeA) = Serializer $ \ state -> let
    (state', f) = computeF state
    in second f $ computeA state'

instance Monad Serializer where
  (>>=) :: Serializer a -> (a -> Serializer b) -> Serializer b
  (>>=) (Serializer computeA) f = Serializer $ \ state -> let
    (state', a) = computeA state
    (Serializer computeB) = f a
    in computeB state'

getSerializerState :: Serializer SerializerState
getSerializerState = Serializer $ \ state -> (state, state)

setSerializerState :: SerializerState -> Serializer ()
setSerializerState state = Serializer $ const (state, ())

modifySerializerState :: (SerializerState -> SerializerState) -> Serializer ()
modifySerializerState f = getSerializerState >>= setSerializerState . f

-- | Append a text to the serializer buffer. Does not check whether the emitted text contains a newline.

emit :: StrictText -> Serializer ()
emit t = Serializer $ \ state -> (state { builder = \ suffix -> t : state.builder suffix }, ())

lineBreak :: Serializer ()
lineBreak = do
  spaces <- indentation <$!> getSerializerState
  emit "\n"
  emit $ StrictText.replicate (fromIntegral spaces) " "

recordField :: (Word -> Serializer ()) -> Serializer b -> Serializer b
recordField labelField dumpField = do
  maybeFieldIndex <- currentFieldIndex <$!> getSerializerState

  -- get current field number
  n <- case maybeFieldIndex of
    Nothing -> emit " {" >> pure 0
    Just n  -> emit "," >> pure n

  -- not on the same line as the previous field
  lineBreak

  -- what is the field called?
  labelField n
  emit ": "

  result <- dumpField

  -- update field count
  modifySerializerState $ \ state -> state { currentFieldIndex = Just $ n + 1 }

  pure result

unnamedField :: Serializer a -> Serializer a
unnamedField = recordField serializer

namedField :: StrictText -> Serializer a -> Serializer a
namedField name = recordField (const $ emit name)

customField :: Serializer () -> Serializer a -> Serializer a
customField = recordField . const

beginFields :: Maybe (Serializer ()) -> Serializer b -> Serializer b
beginFields name body = do
  oldIndentation <- indentation <$!> getSerializerState
  fromMaybe (pure ()) name
  -- increase the indentation
  modifySerializerState $ \ state -> state { currentFieldIndex = Nothing, indentation = state.increaseIndentation state.indentation  }

  -- dump the body
  result <- body

  -- reset the indentation
  modifySerializerState $ \ state -> state { indentation = oldIndentation }

  -- optionally end the braces
  finalFieldIndex <- currentFieldIndex <$!> getSerializerState
  case finalFieldIndex of
    Nothing -> pure ()
    Just _  -> lineBreak >> emit "}"

  pure result

inNamedConstructor :: StrictText -> Serializer a -> Serializer a
inNamedConstructor name = beginFields $ Just (emit "::" >> emit name)


inDatatype :: StrictText -> Serializer b -> Serializer b
inDatatype name body = do
  emit name

  body

inDatatypeOf :: Typeable a => a -> Serializer b -> Serializer b
inDatatypeOf x s = let
  !typeName = StrictText.pack . show . typeOf $ x
  in inDatatype typeName s

run :: Serializer () -> LazyText
run (Serializer computeUnit) = let
  initialState = SerializerState {builder=mempty, indentation=0, increaseIndentation=(+4), currentFieldIndex=Nothing}
  endState = fst (computeUnit initialState)
  in LazyText.fromChunks $ endState.builder []

class PrettySerialize a where
  serialize :: a -> LazyText
  serialize = run . serializer
  serializer :: a -> Serializer ()

newtype ShowPrettySerialize a = ShowPrettySerialize  a

instance Show a => PrettySerialize (ShowPrettySerialize a) where
  serializer :: ShowPrettySerialize a -> Serializer ()
  serializer = emit . fromString . show . (\ (ShowPrettySerialize a) -> a)

-- numbers
deriving via ShowPrettySerialize Int instance PrettySerialize Int
deriving via ShowPrettySerialize Int8 instance PrettySerialize Int8
deriving via ShowPrettySerialize Word instance PrettySerialize Word
deriving via ShowPrettySerialize Bool instance PrettySerialize Bool
deriving via ShowPrettySerialize Float instance PrettySerialize Float
deriving via ShowPrettySerialize Int16 instance PrettySerialize Int16
deriving via ShowPrettySerialize Int32 instance PrettySerialize Int32
deriving via ShowPrettySerialize Int64 instance PrettySerialize Int64
deriving via ShowPrettySerialize Word8 instance PrettySerialize Word8
deriving via ShowPrettySerialize Double instance PrettySerialize Double
deriving via ShowPrettySerialize Word16 instance PrettySerialize Word16
deriving via ShowPrettySerialize Word32 instance PrettySerialize Word32
deriving via ShowPrettySerialize Word64 instance PrettySerialize Word64
deriving via ShowPrettySerialize Integer instance PrettySerialize Integer
deriving via ShowPrettySerialize Natural instance PrettySerialize Natural
deriving via ShowPrettySerialize (Ratio a) instance Show a => PrettySerialize (Ratio a)

-- text
deriving via ShowPrettySerialize LazyText instance PrettySerialize LazyText
deriving via ShowPrettySerialize ByteString instance PrettySerialize ByteString
deriving via ShowPrettySerialize StrictText instance PrettySerialize StrictText
deriving via ShowPrettySerialize LazyByteString instance PrettySerialize LazyByteString

-- tuples
deriving via Generically (Solo a) instance PrettySerialize a => PrettySerialize (Solo a)
deriving via Generically (a, b) instance (PrettySerialize a, PrettySerialize b) => PrettySerialize (a, b)

-- functors
deriving via Generically (Maybe a) instance PrettySerialize a => PrettySerialize (Maybe a)
deriving via Generically (Either l r) instance (PrettySerialize l, PrettySerialize r) => PrettySerialize (Either l r)
deriving via Generically (Identity a) instance (PrettySerialize a) => PrettySerialize (Identity a)
deriving via Generically (Const a b) instance (PrettySerialize a) => PrettySerialize (Const a b)
deriving via Generically (Compose a b c) instance (PrettySerialize (a (b c))) => PrettySerialize (Compose a b c)

-- containers

instance (Typeable a, PrettySerialize a) => PrettySerialize (Set a) where
  serializer :: Set a -> Serializer ()
  serializer set = do
    inDatatypeOf set $ do -- set type name
      beginFields Nothing $ do -- skip constructor name
        namedField "elements" $ do -- emit field
          serializer $ Set.toList set

class ExtractMappings a where
  type KeyType a :: Type
  type ValueType a :: Type
  keyValuePairs :: a -> List (KeyType a, ValueType a)

instance Data.IArray.Ix index => ExtractMappings (Data.IArray.Array index element) where

  type KeyType (Data.IArray.Array index element) = index

  type ValueType (Data.IArray.Array index element) = element

  keyValuePairs :: Data.IArray.Array index element -> [(KeyType (Data.IArray.Array index element), ValueType (Data.IArray.Array index element))]
  keyValuePairs = Data.IArray.assocs

instance ExtractMappings (List a) where
  type KeyType   (List a) = Natural
  type ValueType (List a) = a

  keyValuePairs :: [a] -> [(KeyType [a], ValueType [a])]
  keyValuePairs = zip [0..]

instance ExtractMappings (Map k v) where
  type KeyType   (Map k v) = k
  type ValueType (Map k v) = v

  keyValuePairs :: Map k v -> [(KeyType (Map k v), ValueType (Map k v))]
  keyValuePairs = Map.toList

newtype KeyValueSerialize a = KeyValueSerialize a


instance (ExtractMappings a, Typeable a, PrettySerialize (KeyType a), PrettySerialize (ValueType a)) => PrettySerialize (KeyValueSerialize a) where
  serializer :: KeyValueSerialize a -> Serializer ()
  serializer (KeyValueSerialize mapping) = do
    inDatatypeOf mapping $ do -- set type name
      beginFields Nothing $ do -- don't set constructor name
        forM_ (keyValuePairs mapping) $ \ (key, value) -> do -- emit fields
          customField (serializer key) (serializer value)

deriving via KeyValueSerialize (Data.IArray.Array index element) instance (Data.IArray.Ix index, Typeable index, Typeable element, PrettySerialize index, PrettySerialize element) => PrettySerialize (Data.IArray.Array index element)
deriving via KeyValueSerialize (List a) instance (Typeable a, PrettySerialize a) => PrettySerialize (List a)
deriving via KeyValueSerialize (Map k v) instance (Typeable k, Typeable v, PrettySerialize k, PrettySerialize v) => PrettySerialize (Map k v)


class GenericPrettySerialize self where
  genericDumpSerializer :: self x -> Serializer ()

instance (Generic a, GenericPrettySerialize (Rep a)) => PrettySerialize (Generically a) where
  serializer :: Generically a -> Serializer ()
  serializer (Generically a) = genericDumpSerializer . from $ a

-- unit value
instance GenericPrettySerialize U1 where
  genericDumpSerializer :: U1 x -> Serializer ()
  genericDumpSerializer _ = pure ()

-- a field
instance (PrettySerialize a) => GenericPrettySerialize (K1 i a) where
  genericDumpSerializer :: K1 i a x -> Serializer ()
  genericDumpSerializer (K1 fieldValue) = serializer fieldValue

-- unnamed field meta
instance GenericPrettySerialize a => GenericPrettySerialize (M1 tag (MetaSel Nothing isUnpacked isStrictSource isStrictCompiler) a) where
  genericDumpSerializer :: M1 tag (MetaSel Nothing isUnpacked isStrictSource isStrictCompiler) a x -> Serializer ()
  genericDumpSerializer (M1 f) = unnamedField (genericDumpSerializer f)

-- named field meta
instance (KnownSymbol name, GenericPrettySerialize a) => GenericPrettySerialize (M1 tag (MetaSel (Just name) isUnpacked isStrictSource isStrictCompiler) a) where
  genericDumpSerializer :: M1 tag (MetaSel (Just name) isUnpacked isStrictSource isStrictCompiler) a x -> Serializer ()
  genericDumpSerializer (M1 f) = namedField (fromString fieldName) (genericDumpSerializer f)
    where
      fieldName = symbolVal (Proxy @name)

-- constructor metadata
instance (GenericPrettySerialize a, KnownSymbol name) => GenericPrettySerialize (M1 tag (MetaCons name fixity isRecord) a) where
  genericDumpSerializer :: M1 tag (MetaCons name fixity isRecord) a x -> Serializer ()
  genericDumpSerializer (M1 c) = inNamedConstructor (fromString contructorName) (genericDumpSerializer c)
    where
      contructorName = symbolVal (Proxy @name)

-- datatype metadata
instance (KnownSymbol name, GenericPrettySerialize a) => GenericPrettySerialize (M1 tag (MetaData name module_ package isNewtype) a) where
  genericDumpSerializer :: M1 tag (MetaData name module_ package isNewtype) a x -> Serializer ()
  genericDumpSerializer (M1 d) = inDatatype (fromString dataName) (genericDumpSerializer d)
    where
      dataName = symbolVal (Proxy @name)

-- sum type options
instance (GenericPrettySerialize l, GenericPrettySerialize r) => GenericPrettySerialize (l :+: r) where
  genericDumpSerializer :: (:+:) l r x -> Serializer ()
  genericDumpSerializer = \case
    L1 c -> genericDumpSerializer c
    R1 c -> genericDumpSerializer c

-- product type
instance (GenericPrettySerialize l, GenericPrettySerialize r) => GenericPrettySerialize (l :*: r) where
  genericDumpSerializer :: (:*:) l r x -> Serializer ()
  genericDumpSerializer (l :*: r) = do
    genericDumpSerializer l
    genericDumpSerializer r