{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

import Control.Applicative ((<$>))
import Control.Monad (mzero)
import Data.Aeson
import Data.Aeson.Types
import qualified Data.Aeson as A
import qualified Data.ByteString.Lazy as BL
import Data.Monoid ((<>))
import Data.Proxy (Proxy(Proxy))
import GHC.TypeLits (KnownSymbol, Symbol, symbolVal, sameSymbol)
import Data.Text (pack)
import Data.Type.Equality 

data TypeKeyOf (a :: *) (x :: Symbol) where 
  IntK :: Int `TypeKeyOf` "int" 
  StringK :: String `TypeKeyOf` "string" 

type IsTypeKey a x = (ToJSON a, FromJSON a, KnownSymbol x)

isTypeKey :: TypeKeyOf a x -> (IsTypeKey a x => r) -> r 
isTypeKey IntK k = k 
isTypeKey StringK k = k 

keyOf :: TypeKeyOf a x -> Proxy x 
keyOf _ = Proxy 

instance ToJSON (TypeKeyOf a x) where 
  toJSON k = isTypeKey k (A.String . pack . symbolVal . keyOf $ k)

data SomeTypeKey = forall a x . TK (TypeKeyOf a x) 

instance FromJSON SomeTypeKey where 
    parseJSON (A.String s) 
      | s == "int" = return $ TK IntK 
      | s == "string" = return $ TK StringK 
    parseJSON _      = mzero

data Payload where
  Payload :: a `TypeKeyOf` s -> a -> Payload 

instance ToJSON Payload where
    toJSON (Payload k a) = 
      object [ "type" .= k 
             , isTypeKey k $ "data" .= a
             ]

instance FromJSON Payload where
    parseJSON (Object v) = 
      (v .: "type") >>= \(TK q) -> isTypeKey q (Payload q <$> v .: "data") 
    parseJSON _ = mzero 

-- | Show intance for ghci
instance Show Payload where
    show (Payload k a) = typeKeyShow k $ isTypeKey k $ 
      "Payload " <> symbolVal (keyOf k) <> " " <> show a

typeKeyShow :: TypeKeyOf a x -> (Show a => r) -> r 
typeKeyShow IntK k = k 
typeKeyShow StringK k = k 

jsonString :: BL.ByteString
jsonString = "{\"type\": \"string\", \"data\": \"cool\"}"


class HasTypeKey a (x :: Symbol) | x -> a where 
  typeKey :: TypeKeyOf a x 

instance HasTypeKey Int "int" where typeKey = IntK 
instance HasTypeKey String "string" where typeKey = StringK 

typeKeyOf :: HasTypeKey a x => Proxy x -> TypeKeyOf a x 
typeKeyOf _ = typeKey 

sameKey :: TypeKeyOf a x -> TypeKeyOf a' x' -> Maybe ('(a, x) :~: '(a', x'))
sameKey IntK IntK = Just Refl 
sameKey StringK StringK = Just Refl 
sameKey _ _ = Nothing 

extractPayload :: HasTypeKey a x => Proxy x -> Payload -> Maybe a 
extractPayload t' (Payload t x) = fmap (\Refl -> x) $ sameKey t (typeKeyOf t')