--------------------------------------------------------------------------------
{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}
module Network.WebSockets.Hybi13
    ( headerVersions
    , finishRequest
    , finishResponse
    , encodeMessages
    , decodeMessages
    , createRequest

      -- Internal (used for testing)
    , encodeFrame
    ) where


--------------------------------------------------------------------------------
import qualified Blaze.ByteString.Builder              as B
import           Control.Applicative                   (pure, (<$>))
import           Control.Exception                     (throw)
import           Control.Monad                         (liftM)
import qualified Data.Attoparsec.ByteString            as A
import           Data.Binary.Get                       (getWord16be,
                                                        getWord64be, runGet)
import           Data.Binary.Put                       (putWord16be, runPut)
import           Data.Bits                             ((.&.), (.|.))
import           Data.ByteString                       (ByteString)
import qualified Data.ByteString.Base64                as B64
import           Data.ByteString.Char8                 ()
import qualified Data.ByteString.Lazy                  as BL
import           Data.Digest.Pure.SHA                  (bytestringDigest, sha1)
import           Data.Int                              (Int64)
import           Data.IORef
import           Data.Monoid                           (mappend, mconcat,
                                                        mempty)
import           Data.Tuple                            (swap)
import           System.Entropy                        as R
import           System.Random                         (RandomGen, newStdGen)


--------------------------------------------------------------------------------
import           Network.WebSockets.Http
import           Network.WebSockets.Hybi13.Demultiplex
import           Network.WebSockets.Hybi13.Mask
import           Network.WebSockets.Stream             (Stream)
import qualified Network.WebSockets.Stream             as Stream
import           Network.WebSockets.Types


--------------------------------------------------------------------------------
headerVersions :: [ByteString]
headerVersions = ["13"]


--------------------------------------------------------------------------------
finishRequest :: RequestHead
              -> Headers
              -> Response
finishRequest reqHttp headers =
    let !key     = getRequestHeader reqHttp "Sec-WebSocket-Key"
        !hash    = hashKey key
        !encoded = B64.encode hash
    in response101 (("Sec-WebSocket-Accept", encoded):headers) ""


--------------------------------------------------------------------------------
finishResponse :: RequestHead
               -> ResponseHead
               -> Response
finishResponse request response
    -- Response message should be one of
    --
    -- - WebSocket Protocol Handshake
    -- - Switching Protocols
    --
    -- But we don't check it for now
    | responseCode response /= 101  = throw $ MalformedResponse response
        "Wrong response status or message."
    | responseHash /= challengeHash = throw $ MalformedResponse response
        "Challenge and response hashes do not match."
    | otherwise                     =
        Response response ""
  where
    key           = getRequestHeader  request  "Sec-WebSocket-Key"
    responseHash  = getResponseHeader response "Sec-WebSocket-Accept"
    challengeHash = B64.encode $ hashKey key


--------------------------------------------------------------------------------
encodeMessage :: RandomGen g => ConnectionType -> g -> Message -> (g, B.Builder)
encodeMessage conType gen msg = (gen', builder `mappend` B.flush)
  where
    mkFrame      = Frame True False False False
    (mask, gen') = case conType of
        ServerConnection -> (Nothing, gen)
        ClientConnection -> randomMask gen
    builder      = encodeFrame mask $ case msg of
        (ControlMessage (Close code pl)) -> mkFrame CloseFrame $
            runPut (putWord16be code) `mappend` pl
        (ControlMessage (Ping pl))       -> mkFrame PingFrame   pl
        (ControlMessage (Pong pl))       -> mkFrame PongFrame   pl
        (DataMessage (Text pl))          -> mkFrame TextFrame   pl
        (DataMessage (Binary pl))        -> mkFrame BinaryFrame pl


--------------------------------------------------------------------------------
encodeMessages
    :: ConnectionType
    -> Stream
    -> IO (Message -> IO ())
encodeMessages conType stream = do
    genRef <- newIORef =<< newStdGen
    return $ \msg -> do
        builder <- atomicModifyIORef genRef $ \s -> encodeMessage conType s msg
        Stream.write stream (B.toLazyByteString builder)


--------------------------------------------------------------------------------
encodeFrame :: Mask -> Frame -> B.Builder
encodeFrame mask f = B.fromWord8 byte0 `mappend`
    B.fromWord8 byte1 `mappend` len `mappend` maskbytes `mappend`
    B.fromLazyByteString (maskPayload mask (framePayload f))
  where
    byte0  = fin .|. rsv1 .|. rsv2 .|. rsv3 .|. opcode
    fin    = if frameFin f  then 0x80 else 0x00
    rsv1   = if frameRsv1 f then 0x40 else 0x00
    rsv2   = if frameRsv2 f then 0x20 else 0x00
    rsv3   = if frameRsv3 f then 0x10 else 0x00
    opcode = case frameType f of
        ContinuationFrame -> 0x00
        TextFrame         -> 0x01
        BinaryFrame       -> 0x02
        CloseFrame        -> 0x08
        PingFrame         -> 0x09
        PongFrame         -> 0x0a

    (maskflag, maskbytes) = case mask of
        Nothing -> (0x00, mempty)
        Just m  -> (0x80, B.fromByteString m)

    byte1 = maskflag .|. lenflag
    len'  = BL.length (framePayload f)
    (lenflag, len)
        | len' < 126     = (fromIntegral len', mempty)
        | len' < 0x10000 = (126, B.fromWord16be (fromIntegral len'))
        | otherwise      = (127, B.fromWord64be (fromIntegral len'))


--------------------------------------------------------------------------------
decodeMessages
    :: Stream
    -> IO (IO (Maybe Message))
decodeMessages stream = do
    dmRef <- newIORef emptyDemultiplexState
    return $ go dmRef
  where
    go dmRef = do
        mbFrame <- Stream.parse stream parseFrame
        case mbFrame of
            Nothing    -> return Nothing
            Just frame -> do
                mbMsg <- atomicModifyIORef dmRef $
                    \s -> swap $ demultiplex s frame
                case mbMsg of
                    Nothing  -> go dmRef
                    Just msg -> return (Just msg)


--------------------------------------------------------------------------------
-- | Parse a frame
parseFrame :: A.Parser Frame
parseFrame = do
    byte0 <- A.anyWord8
    let fin    = byte0 .&. 0x80 == 0x80
        rsv1   = byte0 .&. 0x40 == 0x40
        rsv2   = byte0 .&. 0x20 == 0x20
        rsv3   = byte0 .&. 0x10 == 0x10
        opcode = byte0 .&. 0x0f

    ft <- case opcode of
            0x00 -> return ContinuationFrame
            0x01 -> return TextFrame
            0x02 -> return BinaryFrame
            0x08 -> return CloseFrame
            0x09 -> return PingFrame
            0x0a -> return PongFrame
            _    -> fail $ "Unknown opcode: " ++ show opcode

    byte1 <- A.anyWord8
    let mask = byte1 .&. 0x80 == 0x80
        lenflag = fromIntegral (byte1 .&. 0x7f)

    len <- case lenflag of
        126 -> fromIntegral . runGet' getWord16be <$> A.take 2
        127 -> fromIntegral . runGet' getWord64be <$> A.take 8
        _   -> return lenflag

    masker <- maskPayload <$> if mask then Just <$> A.take 4 else pure Nothing

    chunks <- take64 len

    return $ Frame fin rsv1 rsv2 rsv3 ft (masker $ BL.fromChunks chunks)
  where
    runGet' g = runGet g . BL.fromChunks . return

    take64 :: Int64 -> A.Parser [ByteString]
    take64 n
        | n <= 0    = return []
        | otherwise = do
            let n' = min intMax n
            chunk <- A.take (fromIntegral n')
            (chunk :) <$> take64 (n - n')
      where
        intMax :: Int64
        intMax = fromIntegral (maxBound :: Int)


--------------------------------------------------------------------------------
hashKey :: ByteString -> ByteString
hashKey key = unlazy $ bytestringDigest $ sha1 $ lazy $ key `mappend` guid
  where
    guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    lazy = BL.fromChunks . return
    unlazy = mconcat . BL.toChunks


--------------------------------------------------------------------------------
createRequest :: ByteString
              -> ByteString
              -> Bool
              -> Headers
              -> IO RequestHead
createRequest hostname path secure customHeaders = do
    key <- B64.encode `liftM`  getEntropy 16
    return $ RequestHead path (headers key ++ customHeaders) secure
  where
    headers key =
        [ ("Host"                   , hostname     )
        , ("Connection"             , "Upgrade"    )
        , ("Upgrade"                , "websocket"  )
        , ("Sec-WebSocket-Key"      , key          )
        , ("Sec-WebSocket-Version"  , versionNumber)
        ]

    versionNumber = head headerVersions
