{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE StandaloneDeriving     #-}
module Network.RPC where

import Control.Applicative
import qualified Text.ParserCombinators.ReadP as RP
import Data.Digest.CRC32C
import Data.Word
import Data.Monoid
import Data.Hashable
import Data.String
import Data.Bits
import Data.ByteString (ByteString)
import Data.Kind       (Constraint)
import Data.Data
import Data.Default
import Data.List.Split
import Data.Ord
import Data.IP
import Network.Socket
import Text.PrettyPrint as PP hiding ((<>))
import Text.PrettyPrint.HughesPJClass  hiding (($$), (<>))
import Text.Read        (readMaybe)
import Data.Serialize                  as S
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString as BS
import Data.ByteString.Base16          as Base16
import System.Entropy

class (Eq a, Serialize a, Typeable a, Hashable a, Pretty a)
    => Address a where
  toSockAddr   :: a        -> SockAddr
  fromSockAddr :: SockAddr -> Maybe a

fromAddr :: (Address a, Address b) => a -> Maybe b
fromAddr = fromSockAddr . toSockAddr

-- | Note that port is zeroed.
instance Address IPv4 where
  toSockAddr = SockAddrInet 0 . toHostAddress
  fromSockAddr (SockAddrInet _ h) = Just (fromHostAddress h)
  fromSockAddr  _                 = Nothing

-- | Note that port is zeroed.
instance Address IPv6 where
  toSockAddr h = SockAddrInet6 0 0 (toHostAddress6 h) 0
  fromSockAddr (SockAddrInet6 _ _ h _) = Just (fromHostAddress6 h)
  fromSockAddr  _                      = Nothing

-- | Note that port is zeroed.
instance Address IP where
  toSockAddr (IPv4 h) = toSockAddr h
  toSockAddr (IPv6 h) = toSockAddr h
  fromSockAddr sa =
        IPv4 <$> fromSockAddr sa
    <|> IPv6 <$> fromSockAddr sa




data MessageClass = Error | Query | Response
 deriving (Eq,Ord,Enum,Bounded,Data,Show,Read)

class Envelope envelope where
    type TransactionID envelope
    data NodeId envelope

    envelopePayload     :: envelope a -> a
    envelopeTransaction :: envelope a -> TransactionID envelope
    envelopeClass       :: envelope a -> MessageClass

    -- | > buildReply self addr qry response
    --
    --    [ self ]     this node's id.
    --
    --    [ addr ]     SockAddr of query origin.
    --
    --    [ qry ]      received query message.
    --
    --    [ response ] response payload.
    --
    -- Returns: response message envelope
    buildReply :: NodeId envelope -> SockAddr -> envelope a -> b -> envelope b

-- | In Kademlia, the distance metric is XOR and the result is
-- interpreted as an unsigned integer.
newtype NodeDistance nodeid = NodeDistance nodeid
  deriving (Eq, Ord)

-- | distance(A,B) = |A xor B| Smaller values are closer.
distance :: Bits nid => nid -> nid -> NodeDistance nid
distance a b = NodeDistance $ xor a b

instance Serialize nodeid => Show (NodeDistance nodeid) where
  show (NodeDistance w) = Char8.unpack $ Base16.encode $ S.encode w

instance Serialize nodeid => Pretty (NodeDistance nodeid) where
  pPrint n = text $ show n


-- | When 'get'ing an IP it must be 'isolate'd to the appropriate
-- number of bytes since we have no other way of telling which
-- address type we are trying to parse
instance Serialize IP where
    put (IPv4 ip) = put ip
    put (IPv6 ip) = put ip

    get = do
      n <- remaining
      case n of
        4  -> IPv4 <$> get
        16 -> IPv6 <$> get
        _ -> fail (show n ++ " is the wrong number of remaining bytes to parse IP")

instance Serialize IPv4 where
    put = putWord32host    .  toHostAddress
    get = fromHostAddress <$> getWord32host

instance Serialize IPv6 where
    put ip = put $ toHostAddress6 ip
    get = fromHostAddress6 <$> get

instance Pretty IPv4 where
  pPrint = PP.text . show
  {-# INLINE pPrint #-}

instance Pretty IPv6 where
  pPrint = PP.text . show
  {-# INLINE pPrint #-}

instance Pretty IP where
  pPrint = PP.text . show
  {-# INLINE pPrint #-}

instance Hashable IPv4 where
  hashWithSalt = hashUsing toHostAddress
  {-# INLINE hashWithSalt #-}

instance Hashable IPv6 where
  hashWithSalt s a = hashWithSalt s (toHostAddress6 a)

instance Hashable IP where
  hashWithSalt s (IPv4 h) = hashWithSalt s h
  hashWithSalt s (IPv6 h) = hashWithSalt s h





data NodeAddr a = NodeAddr
  { nodeHost ::                !a
  , nodePort :: {-# UNPACK #-} !PortNumber
  } deriving (Eq, Ord, Typeable, Functor, Foldable, Traversable)

instance Show a => Show (NodeAddr a) where
  showsPrec i NodeAddr {..}
    = showsPrec i nodeHost <> showString ":" <> showsPrec i nodePort

instance Read (NodeAddr IPv4) where
  readsPrec i = RP.readP_to_S $ do
    ipv4 <- RP.readS_to_P (readsPrec i)
    _    <- RP.char ':'
    port <- toEnum <$> RP.readS_to_P (readsPrec i)
    return $ NodeAddr ipv4 port

-- | @127.0.0.1:6882@
instance Default (NodeAddr IPv4) where
  def = "127.0.0.1:6882"

-- | KRPC compatible encoding.
instance Serialize a => Serialize (NodeAddr a) where
  get = NodeAddr <$> get <*> get
  {-# INLINE get #-}
  put NodeAddr {..} = put nodeHost >> put nodePort
  {-# INLINE put #-}

-- | Example:
--
--   @nodePort \"127.0.0.1:6881\" == 6881@
--
instance IsString (NodeAddr IPv4) where
  fromString str
    | [hostAddrStr, portStr] <- splitWhen (== ':') str
    , Just hostAddr <- readMaybe hostAddrStr
    , Just portNum  <- toEnum <$> readMaybe portStr
                = NodeAddr hostAddr portNum
    | otherwise = error $ "fromString: unable to parse (NodeAddr IPv4): " ++ str

instance Hashable PortNumber where
  hashWithSalt s = hashWithSalt s . fromEnum
  {-# INLINE hashWithSalt #-}

instance Pretty PortNumber where
  pPrint = PP.int . fromEnum
  {-# INLINE pPrint #-}


instance Hashable a => Hashable (NodeAddr a) where
  hashWithSalt s NodeAddr {..} = hashWithSalt s (nodeHost, nodePort)
  {-# INLINE hashWithSalt #-}

instance Pretty ip => Pretty (NodeAddr ip) where
  pPrint NodeAddr {..} = pPrint nodeHost <> ":" <> pPrint nodePort


instance Serialize PortNumber where
  get = fromIntegral <$> getWord16be
  {-# INLINE get #-}
  put = putWord16be . fromIntegral
  {-# INLINE put #-}




data NodeInfo dht addr u = NodeInfo
  { nodeId   :: !(NodeId dht)
  , nodeAddr :: !(NodeAddr addr)
  , nodeAnnotation :: u
  } deriving (Functor, Foldable, Traversable)

deriving instance ( Show (NodeId dht)
                  , Show addr
                  , Show u ) => Show (NodeInfo dht addr u)

mapAddress :: (addr -> b) -> NodeInfo dht addr u -> NodeInfo dht b u
mapAddress f ni = ni { nodeAddr = fmap f (nodeAddr ni) }

traverseAddress :: Applicative f => (addr -> f b) -> NodeInfo dht addr u -> f (NodeInfo dht b u)
traverseAddress f ni = fmap (\addr -> ni { nodeAddr = addr }) $ traverse f (nodeAddr ni)

-- Warning: Eq and Ord only look at the nodeId field.
instance Eq (NodeId dht) => Eq (NodeInfo dht a u) where
    a == b  = (nodeId a == nodeId b)

instance Ord (NodeId dht) => Ord (NodeInfo dht a u) where
  compare = comparing nodeId


-- TODO WARN is the 'system' random suitable for this?
-- | Generate random NodeID used for the entire session.
--   Distribution of ID's should be as uniform as possible.
--
genNodeId :: forall dht.
             ( Serialize (NodeId dht)
             , FiniteBits (NodeId dht)
             ) => IO (NodeId dht)
genNodeId = either error id . S.decode <$> getEntropy nodeIdSize
 where
    nodeIdSize = finiteBitSize (undefined :: NodeId dht) `div` 8

-- | Generate a random 'NodeId' within a range suitable for a bucket.  To
-- obtain a sample for bucket number /index/ where /is_last/ indicates if this
-- is for the current deepest bucket in our routing table:
--
-- > sample <- genBucketSample nid (bucketRange index is_last)
genBucketSample :: ( FiniteBits (NodeId dht)
                   , Serialize (NodeId dht)
                   ) => NodeId dht -> (Int,Word8,Word8) -> IO (NodeId dht)
genBucketSample n qmb = genBucketSample' getEntropy n qmb

-- | Generalizion of 'genBucketSample' that accepts a byte generator
-- function to use instead of the system entropy.
genBucketSample' :: forall m dht.
                    ( Applicative m
                    , FiniteBits (NodeId dht)
                    , Serialize (NodeId dht)
                    ) =>
    (Int -> m ByteString) -> NodeId dht -> (Int,Word8,Word8) -> m (NodeId dht)
genBucketSample' gen self (q,m,b)
    | q <= 0           =  either error id . S.decode <$> gen nodeIdSize
    | q >= nodeIdSize  =  pure self
    | otherwise        =  either error id . S.decode .  build <$> gen (nodeIdSize - q + 1)
 where
    nodeIdSize = finiteBitSize (undefined :: NodeId dht) `div` 8
    build tl = BS.init hd <> BS.cons (h .|. t) (BS.tail tl)
     where
        hd = BS.take q $ S.encode self
        h = xor b (complement m .&. BS.last hd)
        t = m .&. BS.head tl


class Envelope envelope => WireFormat raw envelope where
    type SerializableTo raw :: * -> Constraint
    type CipherContext raw envelope

    decodeHeaders :: CipherContext raw envelope -> ByteString -> Either String (envelope raw)
    decodePayload :: SerializableTo raw a => envelope raw -> Either String (envelope a)

    encodeHeaders :: CipherContext raw envelope -> envelope raw -> ByteString
    encodePayload :: SerializableTo raw a => envelope a -> envelope raw