{- LANGUAGE ApplicativeDo              -}
{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE ExistentialQuantification  #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TupleSections              #-}
{- LANGUAGE TypeApplications           -}
module Network.Tox.NodeId
    ( NodeInfo(..)
    , NodeId
    , nodeInfo
    , nodeAddr
    , zeroID
    , key2id
    , id2key
    , getIP
    , xorNodeId
    , testNodeIdBit
    , sampleNodeId) where

import Control.Applicative
import Control.Arrow
import Control.Monad
import Crypto.Error.Types                     (CryptoFailable (..),
                                               throwCryptoError)
import Crypto.PubKey.Curve25519
import qualified Data.Aeson                   as JSON
         ;import Data.Aeson                   (FromJSON, ToJSON, (.=))
import Data.Bits.ByteString                   ()
import qualified Data.ByteArray               as BA
         ;import Data.ByteArray               as BA (ByteArrayAccess)
import qualified Data.ByteString              as B
         ;import Data.ByteString              (ByteString)
import qualified Data.ByteString.Base16       as Base16
import qualified Data.ByteString.Base64       as Base64
import qualified Data.ByteString.Char8        as C8
import Data.Char
import Data.Data
import Data.Hashable
import Data.IP
import Data.Serialize                         as S
import Data.Word
import Foreign.Storable
import GHC.TypeLits
import Network.Address                        hiding (nodePort)
import System.IO.Unsafe                       (unsafeDupablePerformIO)
import qualified Text.ParserCombinators.ReadP as RP
import Text.Read
import Data.Bits
import Crypto.Tox
import Foreign.Ptr
import Data.Function
import System.Endian

-- | perform io for hashes that do allocation and ffi.
-- unsafeDupablePerformIO is used when possible as the
-- computation is pure and the output is directly linked
-- to the input. we also do not modify anything after it has
-- been returned to the user.
unsafeDoIO :: IO a -> a
#if __GLASGOW_HASKELL__ > 704
unsafeDoIO = unsafeDupablePerformIO
#else
unsafeDoIO = unsafePerformIO
#endif

unpackPublicKey :: ByteArrayAccess bs => bs -> [Word64]
unpackPublicKey bs = loop 0
  where loop i
            | i == (BA.length bs `div` 8) = []
            | otherwise =
                let !v = unsafeDoIO $ BA.withByteArray bs (\p -> fromBE64 <$> peekElemOff p i)
                 in v : loop (i+1)

packPublicKey :: BA.ByteArray bs => [Word64] -> bs
packPublicKey ws = BA.allocAndFreeze (8 * length ws) $
    flip fix ws $ \loop ys ptr -> case ys of
        [] -> return ()
        x:xs -> do poke ptr (toBE64 x)
                   loop xs (plusPtr ptr 8)

-- We represent the node id redundantly in two formats.  The [Word64] format is
-- convenient for short-circuiting xor/distance comparisons.  The PublicKey
-- format is convenient for encryption.
data NodeId = NodeId [Word64] !(Maybe PublicKey)

instance Eq NodeId where
 (NodeId ws _) == (NodeId xs _)
  =         ws == xs

instance Ord NodeId where
    compare (NodeId ws _) (NodeId xs _) = compare ws xs

instance Sized NodeId where size = ConstSize 32

key2id :: PublicKey -> NodeId
key2id k = NodeId (unpackPublicKey k) (Just k)

bs2id :: ByteString -> NodeId
bs2id bs = uncurry NodeId . (unpackPublicKey &&& Just) $ throwCryptoError . publicKey $ bs

id2key :: NodeId -> PublicKey
id2key (NodeId ws (Just key)) = key
id2key (NodeId key Nothing)   = throwCryptoError . publicKey $ (packPublicKey key :: BA.Bytes)

zeroKey :: PublicKey
zeroKey = throwCryptoError $ publicKey $ B.replicate 32 0

zeroID :: NodeId
zeroID = NodeId (replicate 4 0) (Just zeroKey)

-- | Convert to and from a Base64 variant that uses .- instead of +/.
nmtoken64 :: Bool -> Char -> Char
nmtoken64 False '.' = '+'
nmtoken64 False '-' = '/'
nmtoken64 True  '+' = '.'
nmtoken64 True  '/' = '-'
nmtoken64 _     c   = c

instance Read NodeId where
    readsPrec _ str
        | Right bs <- fmap (BA.drop 1) $ Base64.decode $ C8.pack $ 'A':map (nmtoken64 False) (take 43 str)
        , CryptoPassed pub <- publicKey bs -- B.length bs == 32
                    = [ (key2id pub, drop 43 str) ]
        | otherwise = []

instance Show NodeId where
    show nid = map (nmtoken64 True) $ C8.unpack $ BA.drop 1 $ Base64.encode $ BA.cons 0 $ BA.convert $ id2key nid

instance S.Serialize NodeId where
    get = key2id <$> getPublicKey
    put nid = putPublicKey $ id2key nid

instance Hashable NodeId where
    hashWithSalt salt (NodeId ws _) = hashWithSalt salt (head ws)

testNodeIdBit :: NodeId -> Word -> Bool
testNodeIdBit (NodeId ws _) i -- TODO: Optmize: use ByteArray key if it's available.
  | fromIntegral i < 256 -- 256 bits
  , (q, r) <- quotRem (fromIntegral i) 64
  = testBit (ws !! q) (63 - r)
  |     otherwise      = False

xorNodeId :: NodeId -> NodeId -> NodeId
xorNodeId (NodeId xs _) (NodeId ys _) = NodeId (zipWith xor xs ys) Nothing

sampleNodeId :: Applicative m => (Int -> m ByteString) -> NodeId -> (Int,Word8,Word8) -> m NodeId
sampleNodeId gen (NodeId self k) (q,m,b)
    | q <= 0    = bs2id <$> gen 32
    | q >= 32   = pure (NodeId self k)
    | let (qw,r) = (q+7) `divMod` 8 -- How many Word64 to prepend?
          bw = shiftL (fromIntegral b) (8*(7-r))
          mw = bw - 1 :: Word64
          (hd, t0 : _) = splitAt (qw-1) self
          h = xor bw (complement mw .&. t0)
    = flip fmap (gen $ 8 * (4 - (qw-1)) ) $ \bs ->
        let (w:ws) = unpackPublicKey bs
        in NodeId (hd ++ (h .|. (w .&. mw)) : ws) Nothing

data NodeInfo = NodeInfo
  { nodeId   :: NodeId
  , nodeIP   :: IP
  , nodePort :: PortNumber
  }
 deriving (Eq,Ord)

nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
nodeInfo nid saddr
    | Just ip <- fromSockAddr saddr
    , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port
    | otherwise                       = Left "Address family not supported."


instance ToJSON NodeInfo where
    toJSON (NodeInfo nid (IPv4 ip) port)
        = JSON.object [ "public_key" .= show nid
                      , "ipv4" .= show ip
                      , "port" .= (fromIntegral port :: Int)
                      ]
    toJSON (NodeInfo nid (IPv6 ip6) port)
        | Just ip <- un4map ip6
          = JSON.object [ "public_key" .= show nid
                        , "ipv4" .= show ip
                        , "port" .= (fromIntegral port :: Int)
                        ]
        | otherwise
          = JSON.object [ "public_key" .= show nid
                        , "ipv6" .= show ip6
                        , "port" .= (fromIntegral port :: Int)
                        ]
instance FromJSON NodeInfo where
    parseJSON (JSON.Object v) = do
        nidstr <- v JSON..: "public_key"
        ip6str <- v JSON..:? "ipv6"
        ip4str <- v JSON..:? "ipv4"
        portnum <- v JSON..: "port"
        ip <- maybe empty (return . IPv6) (ip6str >>= readMaybe)
              <|> maybe empty (return . IPv4) (ip4str >>= readMaybe)
        let (bs,_) = Base16.decode (C8.pack nidstr)
            enid   = Base64.decode (C8.pack $ 'A' : map (nmtoken64 False) nidstr)
        idbs <- (guard (B.length bs == 32) >> return bs)
                <|> either fail (return . B.drop 1) enid
        return $ NodeInfo (bs2id idbs) ip (fromIntegral (portnum :: Word16))

getIP :: Word8 -> S.Get IP
getIP 0x02 = IPv4 <$> S.get
getIP 0x0a = IPv6 <$> S.get
getIP 0x82 = IPv4 <$> S.get -- TODO: TCP
getIP 0x8a = IPv6 <$> S.get -- TODO: TCP
getIP x    = fail ("unsupported address family ("++show x++")")

instance Sized NodeInfo where
    size = VarSize $ \(NodeInfo nid ip port) ->
            case ip of
                IPv4 _ -> 39 -- 35 +  4 = 1 + 4 + 2 + 32
                IPv6 _ -> 51 -- 35 + 16 = 1 + 16 + 2 + 32

instance S.Serialize NodeInfo where
    get = do
        addrfam <- S.get :: S.Get Word8
        let fallback = do -- FIXME: Handle unrecognized address families.
                IPv6 <$> S.get
                return $ IPv6 (read "::" :: IPv6)
        ip <- getIP addrfam <|> fallback
        port <- S.get :: S.Get PortNumber
        nid <- S.get
        return $ NodeInfo nid ip port

    put (NodeInfo nid ip port) = do
        case ip of
            IPv4 ip4 -> S.put (2 :: Word8) >> S.put ip4
            IPv6 ip6 -> S.put (10 :: Word8) >> S.put ip6
        S.put port
        S.put nid

hexdigit :: Char -> Bool
hexdigit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'f') || ( 'A' <= c && c <= 'F')

b64digit :: Char -> Bool
b64digit '.' = True
b64digit '+' = True
b64digit '-' = True
b64digit '/' = True
b64digit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'z') || ( 'A' <= c && c <= 'Z')

instance Read NodeInfo where
  readsPrec i = RP.readP_to_S $ do
    RP.skipSpaces
    let n = 43 -- characters in node id.
        parseAddr = RP.between (RP.char '(') (RP.char ')') (RP.munch (/=')'))
                      RP.+++ RP.munch (not . isSpace)
        nodeidAt = do hexhash <- sequence $ replicate n (RP.satisfy b64digit)
                      RP.char '@' RP.+++ RP.satisfy isSpace
                      addrstr <- parseAddr
                      nid <- case Base64.decode $ C8.pack $ 'A' : map (nmtoken64 False) hexhash of
                              Right bs | B.length bs - 1==32 -> return (bs2id $ BA.drop 1 bs)
                              _                              -> fail "Bad node id."
                      return (nid,addrstr)
    (nid,addrstr) <- ( nodeidAt RP.+++ ( (zeroID,) <$> parseAddr) )
    let raddr = do
            ip <- RP.between (RP.char '[') (RP.char ']')
                         (IPv6 <$> RP.readS_to_P (readsPrec i))
                  RP.+++ (IPv4 <$> RP.readS_to_P (readsPrec i))
            _    <- RP.char ':'
            port <- toEnum <$> RP.readS_to_P (readsPrec i)
            return (ip, port)

    (ip,port) <- case RP.readP_to_S raddr addrstr of
                    [] -> fail "Bad address."
                    ((ip,port),_):_ -> return (ip,port)
    return $ NodeInfo nid ip port

-- The Hashable instance depends only on the IP address and port number.
instance Hashable NodeInfo where
  hashWithSalt s ni = hashWithSalt s (nodeIP ni , nodePort ni)
  {-# INLINE hashWithSalt #-}


instance Show NodeInfo where
    showsPrec _ (NodeInfo nid ip port) =
        shows nid . ('@' :) . showsip . (':' :) . shows port
     where
        showsip
            | IPv4 ip4 <- ip                          = shows ip4
            | IPv6 ip6 <- ip , Just ip4 <- un4map ip6 = shows ip4
            | otherwise                               = ('[' :) . shows ip . (']' :)




{-
type NodeId = PubKey

pattern NodeId bs = PubKey bs

-- TODO: This should probably be represented by Curve25519.PublicKey, but
-- ByteString has more instances...
newtype PubKey = PubKey ByteString
 deriving (Eq,Ord,Data, ByteArrayAccess, Bits, Hashable)

instance Serialize PubKey where
    get = PubKey <$> getBytes 32
    put (PubKey bs) = putByteString bs

instance Show PubKey where
    show (PubKey bs) = C8.unpack $ Base16.encode bs

instance FiniteBits PubKey where
    finiteBitSize _ = 256

instance Read PubKey where
    readsPrec _ str
        | (bs, xs) <- Base16.decode $ C8.pack str
        , B.length bs == 32
                    = [ (PubKey bs, drop 64 str) ]
        | otherwise = []




data NodeInfo = NodeInfo
  { nodeId   :: NodeId
  , nodeIP   :: IP
  , nodePort :: PortNumber
  }
 deriving (Eq,Ord,Data)

instance Data PortNumber where
    dataTypeOf _ = mkNoRepType "PortNumber"
    toConstr _     = error "PortNumber.toConstr"
    gunfold _ _    = error "PortNumber.gunfold"

instance ToJSON NodeInfo where
    toJSON (NodeInfo nid (IPv4 ip) port)
        = JSON.object [ "public_key" .= show nid
                      , "ipv4" .= show ip
                      , "port" .= (fromIntegral port :: Int)
                      ]
    toJSON (NodeInfo nid (IPv6 ip6) port)
        | Just ip <- un4map ip6
          = JSON.object [ "public_key" .= show nid
                        , "ipv4" .= show ip
                        , "port" .= (fromIntegral port :: Int)
                        ]
        | otherwise
          = JSON.object [ "public_key" .= show nid
                        , "ipv6" .= show ip6
                        , "port" .= (fromIntegral port :: Int)
                        ]
instance FromJSON NodeInfo where
    parseJSON (JSON.Object v) = do
        nidstr <- v JSON..: "public_key"
        ip6str <- v JSON..:? "ipv6"
        ip4str <- v JSON..:? "ipv4"
        portnum <- v JSON..: "port"
        ip <- maybe empty (return . IPv6) (ip6str >>= readMaybe)
              <|> maybe empty (return . IPv4) (ip4str >>= readMaybe)
        let (bs,_) = Base16.decode (C8.pack nidstr)
        guard (B.length bs == 32)
        return $ NodeInfo (NodeId bs) ip (fromIntegral (portnum :: Word16))

getIP :: Word8 -> S.Get IP
getIP 0x02 = IPv4 <$> S.get
getIP 0x0a = IPv6 <$> S.get
getIP 0x82 = IPv4 <$> S.get -- TODO: TCP
getIP 0x8a = IPv6 <$> S.get -- TODO: TCP
getIP x    = fail ("unsupported address family ("++show x++")")

instance S.Serialize NodeInfo where
    get = do
        addrfam <- S.get :: S.Get Word8
        ip <- getIP addrfam
        port <- S.get :: S.Get PortNumber
        nid <- S.get
        return $ NodeInfo nid ip port

    put (NodeInfo nid ip port) = do
        case ip of
            IPv4 ip4 -> S.put (2 :: Word8) >> S.put ip4
            IPv6 ip6 -> S.put (10 :: Word8) >> S.put ip6
        S.put port
        S.put nid

-- node format:
-- [uint8_t family (2 == IPv4, 10 == IPv6, 130 == TCP IPv4, 138 == TCP IPv6)]
-- [ip (in network byte order), length=4 bytes if ipv4, 16 bytes if ipv6]
-- [port (in network byte order), length=2 bytes]
-- [char array (node_id), length=32 bytes]
--


hexdigit :: Char -> Bool
hexdigit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'f') || ( 'A' <= c && c <= 'F')

instance Read NodeInfo where
  readsPrec i = RP.readP_to_S $ do
    RP.skipSpaces
    let n = 64 -- characters in node id.
        parseAddr = RP.between (RP.char '(') (RP.char ')') (RP.munch (/=')'))
                      RP.+++ RP.munch (not . isSpace)
        nodeidAt = do hexhash <- sequence $ replicate n (RP.satisfy hexdigit)
                      RP.char '@' RP.+++ RP.satisfy isSpace
                      addrstr <- parseAddr
                      nid <- case Base16.decode $ C8.pack hexhash of
                              (bs,_) | B.length bs==32 -> return (PubKey bs)
                              _                        -> fail "Bad node id."
                      return (nid,addrstr)
    (nid,addrstr) <- ( nodeidAt RP.+++ ( (zeroID,) <$> parseAddr) )
    let raddr = do
            ip <- RP.between (RP.char '[') (RP.char ']')
                         (IPv6 <$> RP.readS_to_P (readsPrec i))
                  RP.+++ (IPv4 <$> RP.readS_to_P (readsPrec i))
            _    <- RP.char ':'
            port <- toEnum <$> RP.readS_to_P (readsPrec i)
            return (ip, port)

    (ip,port) <- case RP.readP_to_S raddr addrstr of
                    [] -> fail "Bad address."
                    ((ip,port),_):_ -> return (ip,port)
    return $ NodeInfo nid ip port


-- The Hashable instance depends only on the IP address and port number.
instance Hashable NodeInfo where
  hashWithSalt s ni = hashWithSalt s (nodeIP ni , nodePort ni)
  {-# INLINE hashWithSalt #-}


instance Show NodeInfo where
    showsPrec _ (NodeInfo nid ip port) =
        shows nid . ('@' :) . showsip . (':' :) . shows port
     where
        showsip
            | IPv4 ip4 <- ip                          = shows ip4
            | IPv6 ip6 <- ip , Just ip4 <- un4map ip6 = shows ip4
            | otherwise                               = ('[' :) . shows ip . (']' :)

nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
nodeInfo nid saddr
    | Just ip <- fromSockAddr saddr
    , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port
    | otherwise                       = Left "Address family not supported."

zeroID :: NodeId
zeroID = PubKey $ B.replicate 32 0

-}

nodeAddr :: NodeInfo -> SockAddr
nodeAddr (NodeInfo _ ip port) = setPort port $ toSockAddr ip


newtype ForwardPath (n::Nat) = ForwardPath ByteString
 deriving (Eq, Ord,Data)

{-
class KnownNat n => OnionPacket n where
    mkOnion :: ReturnPath n -> Packet -> Packet
instance OnionPacket 0 where mkOnion _ = id
instance OnionPacket 3 where mkOnion = OnionResponse3
-}