{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Network.DHT.Mainline where

import Data.Digest.CRC32C
import Control.Applicative
import Data.Maybe
import Data.Monoid
import Data.Word
import Data.IP
import Data.BEncode                    as BE
import Data.Bits
import Data.ByteString                 (ByteString)
import Data.ByteString.Base16          as Base16
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString       as BS
import qualified Data.ByteString.Lazy  as L
import Data.Default
import Data.LargeWord
import Data.Serialize                  as S
import Data.String
import Data.Typeable
import Network.KRPC.Message            as KRPC
import Network.RPC as RPC
import Text.PrettyPrint                as PP hiding ((<>))
import Text.PrettyPrint.HughesPJClass  hiding (($$), (<>))

nodeIdSize = finiteBitSize (undefined :: NodeId KMessageOf) `div` 8

instance BEncode (NodeId KMessageOf) where
    toBEncode (NodeId w) = toBEncode $ S.encode w
    fromBEncode bval = fromBEncode bval >>= S.decode

-- instance BEncode NodeId where TODO

-- TODO: put this somewhere appropriate
instance (Serialize a, Serialize b) => Serialize (LargeKey a b) where
   put (LargeKey lo hi) = put hi >> put lo
   get = flip LargeKey <$> get <*> get

instance Serialize (NodeId KMessageOf) where
  get = NodeId <$> get
  {-# INLINE get #-}
  put (NodeId bs) = put bs
  {-# INLINE put #-}

-- | ASCII encoded.
instance IsString (NodeId KMessageOf) where
  fromString str
    | length str == nodeIdSize     = NodeId (either error id $ S.decode (fromString str :: ByteString))
    | length str == 2 * nodeIdSize = NodeId (either error id $ S.decode (fst $ Base16.decode $ fromString str))
    |           otherwise = error "fromString: invalid NodeId length"
  {-# INLINE fromString #-}

-- | Meaningless node id, for testing purposes only.
instance Default (NodeId KMessageOf) where
  def = NodeId 0

-- | base16 encoded.
instance Pretty (NodeId KMessageOf) where
  pPrint (NodeId nid) = PP.text $ Char8.unpack $ Base16.encode $ S.encode nid

-- | KRPC 'compact list' compatible encoding: contact information for
-- nodes is encoded as a 26-byte string. Also known as "Compact node
-- info" the 20-byte Node ID in network byte order has the compact
-- IP-address/port info concatenated to the end.
instance Serialize a => Serialize (NodeInfo KMessageOf a ()) where
  get = (\a b -> NodeInfo a b ()) <$> get <*> get
  put NodeInfo {..} = put nodeId >> put nodeAddr

instance Pretty ip => Pretty (NodeInfo KMessageOf ip ()) where
  pPrint NodeInfo {..} = pPrint nodeId <> "@(" <> pPrint nodeAddr <> ")"

instance Pretty ip => Pretty [NodeInfo KMessageOf ip ()] where
  pPrint = PP.vcat . PP.punctuate "," . map pPrint


-- | Yields all 8 DHT neighborhoods available to you given a particular ip
-- address.
bep42s :: Address a => a -> NodeId KMessageOf -> [NodeId KMessageOf]
bep42s addr (NodeId r) = mapMaybe (bep42 addr) rs
 where
    rs = map (NodeId . change3bits r) [0..7]

-- change3bits :: ByteString -> Word8 -> ByteString
-- change3bits bs n = BS.snoc (BS.init bs) (BS.last bs .&. 0xF8 .|. n)

change3bits :: (Num b, Bits b) => b -> b -> b
change3bits bs n = (bs .&. complement 7) .|. n

-- | Modifies a purely random 'NodeId' to one that is related to a given
-- routable address in accordance with BEP 42.
bep42 :: Address a => a -> NodeId KMessageOf -> Maybe (NodeId KMessageOf)
bep42 addr (NodeId r)
 | Just ip <- fmap S.encode (fromAddr addr :: Maybe IPv4)
              <|> fmap S.encode (fromAddr addr :: Maybe IPv6)
    = genBucketSample' retr (NodeId $ crc $ applyMask ip) (3,0x07,0)
 | otherwise
    = Nothing
 where
    ip4mask = "\x03\x0f\x3f\xff" :: ByteString
    ip6mask = "\x01\x03\x07\x0f\x1f\x3f\x7f\xff" :: ByteString
    nbhood_select = (fromIntegral r :: Word8) .&. 7
    retr n = pure $ BS.drop (nodeIdSize - n) $ S.encode r
    crc = flip shiftL (finiteBitSize (NodeId undefined) - 32) . fromIntegral . crc32c . BS.pack
    applyMask ip = case BS.zipWith (.&.) msk ip of
                    (b:bs) -> (b .|. shiftL nbhood_select 5) : bs
                    bs     -> bs
        where msk | BS.length ip == 4  =  ip4mask
                  | otherwise          =  ip6mask



instance Envelope KMessageOf where
    type TransactionID KMessageOf = KRPC.TransactionId

    -- | Each node has a globally unique identifier known as the \"node
    -- ID.\"
    --
    --   Normally, /this/ node id should be saved between invocations
    --   of the client software.
    newtype NodeId KMessageOf = NodeId Word160
      deriving (Show, Eq, Ord, Typeable, Bits, FiniteBits)

    envelopePayload (Q q) = queryArgs q
    envelopePayload (R r) = respVals r
    envelopePayload (E _) = error "TODO: messagePayload for KError"

    envelopeTransaction (Q q) = queryId q
    envelopeTransaction (R r) = respId r
    envelopeTransaction (E e) = errorId e

    envelopeClass (Q _) = Query
    envelopeClass (R _) = Response
    envelopeClass (E _) = Error

    buildReply self addr qry response =
        (R (KResponse response (envelopeTransaction qry) (Just $ ReflectedIP addr)))

instance WireFormat BValue KMessageOf where
    type SerializableTo BValue = BEncode
    type CipherContext BValue KMessageOf = ()

    decodeHeaders _ bs = BE.decode bs >>= BE.fromBEncode
    decodePayload kmsg = mapM BE.fromBEncode kmsg

    encodeHeaders _ kmsg = L.toStrict $ BE.encode kmsg
    encodePayload msg = fmap BE.toBEncode msg