{-
    A parser and encoder for the BitTorrent wire protocol using the
    cereal package.
-}


module Protocol.Wire
    ( Message(..)
    , encodePacket
    , decodeMsg
    , constructBitField
    -- Handshaking
    , initiateHandshake
    , receiveHandshake
    -- Tests
    , testSuite
    )
where

import Control.Applicative hiding (empty)
import Control.Monad

import Data.Monoid
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L

import Data.Serialize
import Data.Serialize.Put
import Data.Serialize.Get

import Data.Char
import System.IO
import System.Log.Logger

import Test.Framework
import Test.Framework.Providers.QuickCheck2
import Test.QuickCheck

import Torrent

------------------------------------------------------------

type BitField    = L.ByteString
type PieceLength = Int

data Message = KeepAlive
             | Choke
             | Unchoke
             | Interested
             | NotInterested
             | Have PieceNum -- Int
             | BitField BitField
             | Request PieceNum Block
             | Piece PieceNum Int B.ByteString
             | Cancel PieceNum Block
             | Port Integer
  deriving (Eq, Show)

instance Arbitrary Message where
    arbitrary = oneof [return KeepAlive, return Choke, return Unchoke, return Interested,
                       return NotInterested,
                       Have <$> arbitrary,
                       BitField <$> arbitrary,
                       Request <$> arbitrary <*> arbitrary,
                       Piece <$> arbitrary <*> arbitrary <*> arbitrary,
                       Cancel <$> arbitrary <*> arbitrary,
                       Port <$> choose (0,16383)]


data HandShake = HandShake 
                  String -- Protocol Header
                  Word64 -- Extension Bias

-- | The Protocol header for the Peer Wire Protocol
protocolHeader :: String
protocolHeader = "BitTorrent protocol"

extensionBasis :: Word64
extensionBasis = 0

extensionFast :: Word64
extensionFast = 4

p8 :: Word8 -> Put
p8 = putWord8

p32be :: Integral a => a -> Put
p32be = putWord32be . fromIntegral

decodeMsg :: Get Message
decodeMsg = get

encodePacket :: Message -> B.ByteString
encodePacket m = mconcat [szEnc, mEnc]
  where mEnc  = encode m
        sz    = B.length mEnc
        szEnc = runPut . p32be $ sz

instance Serialize Message where
    put KeepAlive       = return ()
    put Choke           = p8 0
    put Unchoke         = p8 1
    put Interested      = p8 2
    put NotInterested   = p8 3
    put (Have pn)       = p8 4 *> p32be pn
    put (BitField bf)   = p8 5 *> putLazyByteString bf
    put (Request pn (Block os sz))
                        = p8 6 *> mapM_ p32be [pn,os,sz]
    put (Piece pn os c) = p8 7 *> mapM_ p32be [pn,os] *> putByteString c
    put (Cancel pn (Block os sz))
                        = p8 8 *> mapM_ p32be [pn,os,sz]
    put (Port p)        = p8 9 *> (putWord16be . fromIntegral $ p)

    get =  getKA      <|> getChoke
       <|> getUnchoke <|> getIntr
       <|> getNI      <|> getHave
       <|> getBF      <|> getReq
       <|> getPiece   <|> getCancel
       <|> getPort

getChoke   = byte 0 *> return Choke
getUnchoke = byte 1 *> return Unchoke
getIntr    = byte 2 *> return Interested
getNI      = byte 3 *> return NotInterested
getHave    = byte 4 *> (Have <$> gw32)
getBF      = byte 5 *> (BitField <$> (remaining >>= getLazyByteString . fromIntegral))
getReq     = byte 6 *> (Request  <$> gw32 <*> (Block <$> gw32 <*> gw32))
getPiece   = byte 7 *> (Piece    <$> gw32 <*> gw32 <*> (remaining >>= getByteString))
getCancel  = byte 8 *> (Cancel   <$> gw32 <*> (Block <$> gw32 <*> gw32))
getPort    = byte 9 *> (Port . fromIntegral <$> getWord16be)
getKA      = do
    empty <- isEmpty
    if empty
        then return KeepAlive
        else fail "Non empty message - not a KeepAlive"

gw32 :: Integral a => Get a
gw32 = fromIntegral <$> getWord32be

byte :: Word8 -> Get Word8
byte w = do
    x <- lookAhead getWord8
    if x == w
        then getWord8
        else fail $ "Expected byte: " ++ show w ++ " got: " ++ show x

-- | Size of the protocol header
protocolHeaderSize :: Int
protocolHeaderSize = length protocolHeader

-- | Protocol handshake code. This encodes the protocol handshake part
protocolHandshake :: L.ByteString
protocolHandshake = L.fromChunks . map runPut $
                    [putWord8 $ fromIntegral protocolHeaderSize,
                     mapM_ (putWord8 . fromIntegral . ord) protocolHeader,
                     putWord64be extensionBasis]

toBS :: String -> B.ByteString
toBS = B.pack . map toW8

toLBS :: String -> L.ByteString
toLBS = L.pack . map toW8

fromLBS :: L.ByteString -> String
fromLBS = map (chr . fromIntegral) . L.unpack

toW8 :: Char -> Word8
toW8 = fromIntegral . ord


-- | Receive the header parts from the other end
receiveHeader :: Handle -> Int -> (InfoHash -> Bool)
              -> IO (Either String ([Capabilities], L.ByteString))
receiveHeader h sz ihTst = parseHeader `fmap` B.hGet h sz
  where parseHeader = runGet (headerParser ihTst)

headerParser :: (InfoHash -> Bool) -> Get ([Capabilities], L.ByteString)
headerParser ihTst = do
    hdSz <- getWord8
    when (fromIntegral hdSz /= protocolHeaderSize) $ fail "Wrong header size"
    protoString <- getByteString protocolHeaderSize
    when (protoString /= toBS protocolHeader) $ fail "Wrong protocol header"
    caps <- getWord64be
    ihR  <- liftM fromLBS $ getLazyByteString 20
    unless (ihTst ihR) $ fail "Wrong InfoHash"
    pid <- getLazyByteString 20
    return (decodeCapabilities caps, pid)


data Capabilities = Fast
decodeCapabilities :: Word64 -> [Capabilities]
decodeCapabilities _ = []

-- | Initiate a handshake on a socket
initiateHandshake :: Handle -> PeerId -> InfoHash
                  -> IO (Either String ([Capabilities], L.ByteString))
initiateHandshake handle peerid infohash = do
    debugM "Protocol.Wire" "Sending off handshake message"
    L.hPut handle msg
    hFlush handle
    debugM "Protocol.Wire" "Receiving handshake from other end"
    receiveHeader handle sz (== infohash) -- TODO: Exceptions ?
  where msg = handShakeMessage peerid infohash
        sz = fromIntegral (L.length msg)

-- | Construct a default handshake message from a PeerId and an InfoHash
handShakeMessage :: PeerId -> InfoHash -> L.ByteString
handShakeMessage pid ih =
    L.fromChunks . map runPut $ [putLazyByteString protocolHandshake,
                                 putLazyByteString $ toLBS ih,
                                 putByteString . toBS $ pid]

-- | Receive a handshake on a socket
receiveHandshake :: Handle -> PeerId -> (InfoHash -> Bool) -> InfoHash
                 -> IO (Either String ([Capabilities], L.ByteString))
receiveHandshake h pid ihTst ih = do
    debugM "Protocol.Wire" "Receiving handshake from other end"
    r <- receiveHeader h sz ihTst -- TODO: Exceptions ?
    case r of
        Left err -> return $ Left err
        Right (caps, rpid) ->
            do debugM "Protocol.Wire" "Sending back handshake message"
               L.hPut h msg
               hFlush h
               return $ Right (caps, rpid)
  where msg = handShakeMessage pid ih
        sz = fromIntegral (L.length msg)


-- | The call @constructBitField pieces@ will return the a ByteString suitable for inclusion in a
--   BITFIELD message to a peer.
constructBitField :: Int -> [PieceNum] -> L.ByteString
constructBitField sz pieces = L.pack . build $ m
    where m = map (`elem` pieces) [0..sz-1 + pad]
          pad = 8 - (sz `mod` 8)
          build [] = []
          build l = let (first, rest) = splitAt 8 l
                    in if length first /= 8
                       then error "Wront bitfield"
                       else bytify first : build rest
          bytify [b7,b6,b5,b4,b3,b2,b1,b0] = sum [if b0 then 1   else 0,
                                                  if b1 then 2   else 0,
                                                  if b2 then 4   else 0,
                                                  if b3 then 8   else 0,
                                                  if b4 then 16  else 0,
                                                  if b5 then 32  else 0,
                                                  if b6 then 64  else 0,
                                                  if b7 then 128 else 0]

--
-- -- TESTS

testSuite = testGroup "Protocol/Wire"
  [ testProperty "QC encode-decode/id" propEncodeDecodeId]


propEncodeDecodeId :: Message -> Bool
propEncodeDecodeId m =
    let encoded = encode m
        decoded = decode encoded
    in
        Right m == decoded