content: module Network.QUIC.Packet.Frame (
      encodeFrames,
      encodeFramesWithPadding,
      decodeFramesBS,
      decodeFramesBuffer,
      countZero, -- testing
  ) where
  
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.Short as Short
  import Foreign.Ptr (Ptr, alignPtr, castPtr, minusPtr, plusPtr)
  import Foreign.Storable (alignment, peek)
  import Network.Socket.Internal (zeroMemory)
  
  import Network.QUIC.Imports
  import Network.QUIC.Types
  
  ----------------------------------------------------------------
  
  encodeFrames :: [Frame] -> IO ByteString
  encodeFrames frames = withWriteBuffer 2048 $ \wbuf ->
      mapM_ (encodeFrame wbuf) frames
  
  encodeFramesWithPadding
      :: Buffer
      -> BufferSize
      -> [Frame]
      -> IO Int
      -- ^ payload size without paddings
  encodeFramesWithPadding buf siz frames = do
      zeroMemory buf $ fromIntegral siz -- padding
      wbuf <- newWriteBuffer buf siz
      save wbuf
      mapM_ (encodeFrame wbuf) frames
      savingSize wbuf
  
  encodeFrame :: WriteBuffer -> Frame -> IO ()
  encodeFrame wbuf (Padding n) = replicateM_ n $ write8 wbuf 0x00
  encodeFrame wbuf Ping = write8 wbuf 0x01
  encodeFrame wbuf (Ack (AckInfo largest range1 ranges) (Milliseconds delay)) = do
      write8 wbuf 0x02
      encodeInt' wbuf $ fromIntegral largest
      encodeInt' wbuf $ fromIntegral delay
      encodeInt' wbuf $ fromIntegral $ length ranges
      encodeInt' wbuf $ fromIntegral range1
      mapM_ putRanges ranges
    where
      putRanges (gap, rng) = do
          encodeInt' wbuf $ fromIntegral gap
          encodeInt' wbuf $ fromIntegral rng
  encodeFrame wbuf (ResetStream sid (ApplicationProtocolError err) finalLen) = do
      write8 wbuf 0x04
      encodeInt' wbuf $ fromIntegral sid
      encodeInt' wbuf $ fromIntegral err
      encodeInt' wbuf $ fromIntegral finalLen
  encodeFrame wbuf (StopSending sid (ApplicationProtocolError err)) = do
      write8 wbuf 0x05
      encodeInt' wbuf $ fromIntegral sid
      encodeInt' wbuf $ fromIntegral err
  encodeFrame wbuf (CryptoF off cdata) = do
      write8 wbuf 0x06
      encodeInt' wbuf $ fromIntegral off
      encodeInt' wbuf $ fromIntegral $ BS.length cdata
      copyByteString wbuf cdata
  encodeFrame wbuf (NewToken token) = do
      write8 wbuf 0x07
      encodeInt' wbuf $ fromIntegral $ BS.length token
      copyByteString wbuf token
  encodeFrame wbuf (StreamF sid off dats fin) = do
      let flag0 = 0x08 .|. 0x02 -- len
          flag1
              | off /= 0 = flag0 .|. 0x04 -- off
              | otherwise = flag0
          flag2
              | fin = flag1 .|. 0x01 -- fin
              | otherwise = flag1
      write8 wbuf flag2
      encodeInt' wbuf $ fromIntegral sid
      when (off /= 0) $ encodeInt' wbuf $ fromIntegral off
      encodeInt' wbuf $ fromIntegral $ totalLen dats
      mapM_ (copyByteString wbuf) dats
  encodeFrame wbuf (MaxData n) = do
      write8 wbuf 0x10
      encodeInt' wbuf $ fromIntegral n
  encodeFrame wbuf (MaxStreamData sid n) = do
      write8 wbuf 0x11
      encodeInt' wbuf $ fromIntegral sid
      encodeInt' wbuf $ fromIntegral n
  encodeFrame wbuf (MaxStreams dir ms) = do
      case dir of
          Bidirectional -> write8 wbuf 0x12
          Unidirectional -> write8 wbuf 0x13
      encodeInt' wbuf $ fromIntegral ms
  encodeFrame wbuf (DataBlocked n) = do
      write8 wbuf 0x14
      encodeInt' wbuf $ fromIntegral n
  encodeFrame wbuf (StreamDataBlocked sid n) = do
      write8 wbuf 0x15
      encodeInt' wbuf $ fromIntegral sid
      encodeInt' wbuf $ fromIntegral n
  encodeFrame wbuf (StreamsBlocked dir ms) = do
      case dir of
          Bidirectional -> write8 wbuf 0x16
          Unidirectional -> write8 wbuf 0x17
      encodeInt' wbuf $ fromIntegral ms
  encodeFrame wbuf (NewConnectionID cidInfo retirePriorTo) = do
      write8 wbuf 0x18
      encodeInt' wbuf $ fromIntegral $ cidInfoSeq cidInfo
      encodeInt' wbuf $ fromIntegral retirePriorTo
      let (cid, len) = unpackCID $ cidInfoCID cidInfo
      write8 wbuf len
      copyShortByteString wbuf cid
      let StatelessResetToken token = cidInfoSRT cidInfo
      copyShortByteString wbuf token
  encodeFrame wbuf (RetireConnectionID seqNum) = do
      write8 wbuf 0x19
      encodeInt' wbuf $ fromIntegral seqNum
  encodeFrame wbuf (PathChallenge (PathData pdata)) = do
      write8 wbuf 0x1a
      copyByteString wbuf $ Short.fromShort pdata
  encodeFrame wbuf (PathResponse (PathData pdata)) = do
      write8 wbuf 0x1b
      copyByteString wbuf $ Short.fromShort pdata
  encodeFrame wbuf (ConnectionClose (TransportError err) ftyp reason) = do
      write8 wbuf 0x1c
      encodeInt' wbuf $ fromIntegral err
      encodeInt' wbuf $ fromIntegral ftyp
      encodeInt' wbuf $ fromIntegral $ Short.length reason
      copyShortByteString wbuf reason
  encodeFrame wbuf (ConnectionCloseApp (ApplicationProtocolError err) reason) = do
      write8 wbuf 0x1d
      encodeInt' wbuf $ fromIntegral err
      encodeInt' wbuf $ fromIntegral $ Short.length reason
      copyShortByteString wbuf reason
  encodeFrame wbuf HandshakeDone =
      write8 wbuf 0x1e
  encodeFrame wbuf (UnknownFrame typ) =
      write8 wbuf $ fromIntegral typ
  
  ----------------------------------------------------------------
  
  decodeFramesBS :: ByteString -> IO (Maybe [Frame])
  decodeFramesBS bs = withReadBuffer bs decodeFrames
  
  decodeFramesBuffer :: Buffer -> BufferSize -> IO (Maybe [Frame])
  decodeFramesBuffer buf bufsiz = newReadBuffer buf bufsiz >>= decodeFrames
  
  decodeFrames :: ReadBuffer -> IO (Maybe [Frame])
  decodeFrames rbuf = loop id
    where
      loop frames = do
          ok <- (>= 1) <$> remainingSize rbuf
          if ok
              then do
                  frame <- decodeFrame rbuf
                  case frame of
                      UnknownFrame _ -> return Nothing
                      _ -> loop (frames . (frame :))
              else return $ Just $ frames []
  
  decodeFrame :: ReadBuffer -> IO Frame
  decodeFrame rbuf = do
      ftyp <- fromIntegral <$> decodeInt' rbuf
      case ftyp :: FrameType of
          0x00 -> decodePadding rbuf
          0x01 -> return Ping
          0x02 -> decodeAck rbuf
          -- 0x03 -> Ack with ECN Counts
          0x04 -> decodeResetStream rbuf
          0x05 -> decodeStopSending rbuf
          0x06 -> decodeCrypto rbuf
          0x07 -> decodeNewToken rbuf
          x | 0x08 <= x && x <= 0x0f -> do
              let off = testBit x 2
                  len = testBit x 1
                  fin = testBit x 0
              decodeStream rbuf off len fin
          0x10 -> decodeMaxData rbuf
          0x11 -> decodeMaxStreamData rbuf
          0x12 -> decodeMaxStreams rbuf Bidirectional
          0x13 -> decodeMaxStreams rbuf Unidirectional
          0x14 -> decodeDataBlocked rbuf
          0x15 -> decodeStreamDataBlocked rbuf
          0x16 -> decodeStreamsBlocked rbuf Bidirectional
          0x17 -> decodeStreamsBlocked rbuf Unidirectional
          0x18 -> decodeNewConnectionID rbuf
          0x19 -> decodeRetireConnectionID rbuf
          0x1a -> decodePathChallenge rbuf
          0x1b -> decodePathResponse rbuf
          0x1c -> decodeConnectionClose rbuf
          0x1d -> decodeConnectionCloseApp rbuf
          0x1e -> return HandshakeDone
          x -> return $ UnknownFrame x
  
  decodePadding :: ReadBuffer -> IO Frame
  decodePadding rbuf = do
      n <- withCurrentOffSet rbuf $ \beg -> do
          rest <- remainingSize rbuf
          let end = beg `plusPtr` rest
          countZero beg end
      ff rbuf n
      return $ Padding (n + 1)
  
  countZero :: Ptr Word8 -> Ptr Word8 -> IO Int
  countZero beg0 end0
      | (end0 `minusPtr` beg0) <= ali = fst <$> countBy1 beg0 end0 0
      | otherwise = do
          let beg1 = alignPtr beg0 ali
              end1' = alignPtr end0 ali
              end1
                  | end0 == end1' = end1'
                  | otherwise = end1' `plusPtr` negate ali
          (n1, cont1) <- countBy1 beg0 beg1 0
          if not cont1
              then return n1
              else do
                  (n2, beg2) <- countBy8 (castPtr beg1) (castPtr end1) 0
                  (n3, _) <- countBy1 (castPtr beg2) end0 0
                  return (n1 + n2 + n3)
    where
      ali = alignment (0 :: Word64)
      countBy1 :: Ptr Word8 -> Ptr Word8 -> Int -> IO (Int, Bool)
      countBy1 beg end n
          | beg < end = do
              ftyp <- peek beg
              if ftyp == 0
                  then countBy1 (beg `plusPtr` 1) end (n + 1)
                  else return (n, False)
          | otherwise = return (n, True)
      countBy8 :: Ptr Word64 -> Ptr Word64 -> Int -> IO (Int, Ptr Word64)
      countBy8 beg end n
          | beg < end = do
              ftyp <- peek beg
              if ftyp == 0
                  then countBy8 (beg `plusPtr` ali) end (n + ali)
                  else return (n, beg)
          | otherwise = return (n, beg)
  
  decodeCrypto :: ReadBuffer -> IO Frame
  decodeCrypto rbuf = do
      off <- fromIntegral <$> decodeInt' rbuf
      len <- fromIntegral <$> decodeInt' rbuf
      cdata <- extractByteString rbuf len
      return $ CryptoF off cdata
  
  decodeAck :: ReadBuffer -> IO Frame
  decodeAck rbuf = do
      largest <- fromIntegral <$> decodeInt' rbuf
      delay <- fromIntegral <$> decodeInt' rbuf
      count <- fromIntegral <$> decodeInt' rbuf
      range1 <- fromIntegral <$> decodeInt' rbuf
      ranges <- getRanges count id
      return $ Ack (AckInfo largest range1 ranges) $ Milliseconds delay
    where
      getRanges 0 build = return $ build []
      getRanges n build = do
          gap <- fromIntegral <$> decodeInt' rbuf
          rng <- fromIntegral <$> decodeInt' rbuf
          let n' = n - 1 :: Int
          getRanges n' (build . ((gap, rng) :))
  
  decodeResetStream :: ReadBuffer -> IO Frame
  decodeResetStream rbuf = do
      sID <- fromIntegral <$> decodeInt' rbuf
      err <- ApplicationProtocolError . fromIntegral <$> decodeInt' rbuf
      finalLen <- fromIntegral <$> decodeInt' rbuf
      return $ ResetStream sID err finalLen
  
  decodeStopSending :: ReadBuffer -> IO Frame
  decodeStopSending rbuf = do
      sID <- fromIntegral <$> decodeInt' rbuf
      err <- ApplicationProtocolError . fromIntegral <$> decodeInt' rbuf
      return $ StopSending sID err
  
  decodeNewToken :: ReadBuffer -> IO Frame
  decodeNewToken rbuf = do
      len <- fromIntegral <$> decodeInt' rbuf
      NewToken <$> extractByteString rbuf len
  
  decodeStream :: ReadBuffer -> Bool -> Bool -> Bool -> IO Frame
  decodeStream rbuf hasOff hasLen fin = do
      sID <- fromIntegral <$> decodeInt' rbuf
      off <-
          if hasOff
              then fromIntegral <$> decodeInt' rbuf
              else return 0
      dat <-
          if hasLen
              then do
                  len <- fromIntegral <$> decodeInt' rbuf
                  extractByteString rbuf len
              else do
                  len <- remainingSize rbuf
                  extractByteString rbuf len
      return $ StreamF sID off [dat] fin
  
  ----------------------------------------------------------------
  
  decodeMaxData :: ReadBuffer -> IO Frame
  decodeMaxData rbuf = MaxData . fromIntegral <$> decodeInt' rbuf
  
  decodeMaxStreamData :: ReadBuffer -> IO Frame
  decodeMaxStreamData rbuf = do
      sID <- fromIntegral <$> decodeInt' rbuf
      maxstrdata <- fromIntegral <$> decodeInt' rbuf
      return $ MaxStreamData sID maxstrdata
  
  decodeMaxStreams :: ReadBuffer -> Direction -> IO Frame
  decodeMaxStreams rbuf dir = MaxStreams dir . fromIntegral <$> decodeInt' rbuf
  
  ----------------------------------------------------------------
  
  decodeDataBlocked :: ReadBuffer -> IO Frame
  decodeDataBlocked rbuf = DataBlocked . fromIntegral <$> decodeInt' rbuf
  
  decodeStreamDataBlocked :: ReadBuffer -> IO Frame
  decodeStreamDataBlocked rbuf = do
      sID <- fromIntegral <$> decodeInt' rbuf
      msd <- fromIntegral <$> decodeInt' rbuf
      return $ StreamDataBlocked sID msd
  
  decodeStreamsBlocked :: ReadBuffer -> Direction -> IO Frame
  decodeStreamsBlocked rbuf dir = StreamsBlocked dir . fromIntegral <$> decodeInt' rbuf
  
  ----------------------------------------------------------------
  
  decodeConnectionClose :: ReadBuffer -> IO Frame
  decodeConnectionClose rbuf = do
      err <- TransportError . fromIntegral <$> decodeInt' rbuf
      ftyp <- fromIntegral <$> decodeInt' rbuf
      len <- fromIntegral <$> decodeInt' rbuf
      reason <- extractShortByteString rbuf len
      return $ ConnectionClose err ftyp reason
  
  decodeConnectionCloseApp :: ReadBuffer -> IO Frame
  decodeConnectionCloseApp rbuf = do
      err <- ApplicationProtocolError . fromIntegral <$> decodeInt' rbuf
      len <- fromIntegral <$> decodeInt' rbuf
      reason <- extractShortByteString rbuf len
      return $ ConnectionCloseApp err reason
  
  decodeNewConnectionID :: ReadBuffer -> IO Frame
  decodeNewConnectionID rbuf = do
      seqNum <- fromIntegral <$> decodeInt' rbuf
      retirePriorTo <- fromIntegral <$> decodeInt' rbuf
      cidLen <- fromIntegral <$> read8 rbuf
      cID <- makeCID <$> extractShortByteString rbuf cidLen
      token <- StatelessResetToken <$> extractShortByteString rbuf 16
      return $ NewConnectionID (newCIDInfo seqNum cID token) retirePriorTo
  
  decodeRetireConnectionID :: ReadBuffer -> IO Frame
  decodeRetireConnectionID rbuf = do
      seqNum <- fromIntegral <$> decodeInt' rbuf
      return $ RetireConnectionID seqNum
  
  decodePathChallenge :: ReadBuffer -> IO Frame
  decodePathChallenge rbuf =
      PathChallenge . PathData <$> extractShortByteString rbuf 8
  
  decodePathResponse :: ReadBuffer -> IO Frame
  decodePathResponse rbuf =
      PathResponse . PathData <$> extractShortByteString rbuf 8
  