{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}

module Raw (
  newTraceIdFromHeader,
  newSpanIdFromHeader,
  newHeaderFromTraceId,
  newHeaderFromSpanId,
  showWord64BS,
  readWord64BS,
  asciiWord8ToWord8,
  word8ToAsciiWord8,
) where

import Control.Monad.ST (ST, runST)
import Data.Bits (Bits (complement, shift, (.&.)))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BI
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short.Internal as SBI
import qualified Data.Char as C
import Data.Primitive.ByteArray (
  ByteArray (ByteArray),
  MutableByteArray,
  freezeByteArray,
  indexByteArray,
  newByteArray,
  writeByteArray,
 )
import Data.Primitive.Ptr (writeOffPtr)
import Data.Word (Word64, Word8)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Storable (peekElemOff)
import System.IO.Unsafe (unsafeDupablePerformIO)


newTraceIdFromHeader
  :: ByteString
  -- ^ ASCII numeric text
  -> ShortByteString
newTraceIdFromHeader bs =
  let len = 16 :: Int
      !(ByteArray ba) =
        runST $ do
          mba <- newByteArray len
          let w64 = readWord64BS bs
          writeByteArray mba 0 (0 :: Word64) -- fill zeros to one upper Word64-size area
          writeByteArrayNbo mba 1 w64 -- offset one Word64-size
          freezeByteArray mba 0 len
  in SBI.SBS ba


newSpanIdFromHeader
  :: ByteString
  -- ^ ASCII numeric text
  -> ShortByteString
newSpanIdFromHeader bs =
  let len = 8 :: Int
      !(ByteArray ba) =
        runST $ do
          mba <- newByteArray len
          let w64 = readWord64BS bs
          writeByteArrayNbo mba 0 w64
          freezeByteArray mba 0 len
  in SBI.SBS ba


{- | Write a primitive value to the byte array with network-byte-order (big-endian).
The offset is given in elements of type @a@ rather than in bytes.
-}
writeByteArrayNbo :: MutableByteArray s -> Int -> Word64 -> ST s ()
writeByteArrayNbo mba offset value = do
  writeByteArray mba offset (0 :: Word64)
  loop 0 value
  where
    loop _ 0 = pure ()
    loop 8 _ = pure ()
    loop n v = do
      let
        -- equivelent:
        --   (p, q) = v `divMod` (2 ^ (8 :: Int))
        p = shift v (-8)
        q = v .&. complement (shift p 8)
      writeByteArray mba (8 * (offset + 1) - n - 1) (fromIntegral q :: Word8)
      loop (n + 1) p


readWord64BS :: ByteString -> Word64
readWord64BS (BI.PS fptr _ len) =
  unsafeDupablePerformIO $
    withForeignPtr fptr readWord64Ptr
  where
    readWord64Ptr ptr =
      readWord64PtrOffset 0 0
      where
        readWord64PtrOffset offset acc
          | offset < len = do
              b <- peekElemOff ptr offset
              let n = fromIntegral $ asciiWord8ToWord8 b :: Word64
              readWord64PtrOffset (offset + 1) $ n + acc * 10
          | otherwise = pure acc


asciiWord8ToWord8 :: Word8 -> Word8
asciiWord8ToWord8 b = b - fromIntegral (C.ord '0')


newHeaderFromTraceId :: ShortByteString -> ByteString
newHeaderFromTraceId (SBI.SBS ba) =
  let w64 = indexByteArrayNbo (ByteArray ba) 1
  in showWord64BS w64


newHeaderFromSpanId :: ShortByteString -> ByteString
newHeaderFromSpanId (SBI.SBS ba) =
  let w64 = indexByteArrayNbo (ByteArray ba) 0
  in showWord64BS w64


indexByteArrayNbo :: ByteArray -> Int -> Word64
indexByteArrayNbo ba offset =
  loop 0 0
  where
    loop 8 acc = acc
    loop n acc = loop (n + 1) $ shift acc 8 + word8ToWord64 (indexByteArray ba $ 8 * offset + n)


showWord64BS :: Word64 -> ByteString
showWord64BS value =
  unsafeDupablePerformIO $
    BI.createUptoN 20 writeWord64Ptr -- 20 = length (show (maxBound :: Word64))
  where
    writeWord64Ptr ptr =
      loop (19 :: Int) value 0 False
      where
        loop 0 v offset _ = do
          writeOffPtr ptr offset (word8ToAsciiWord8 $ fromIntegral v)
          pure $ offset + 1
        loop n v offset upper = do
          let (p, q) = v `divMod` (10 ^ n)
          if p == 0 && not upper
            then loop (n - 1) q offset upper
            else do
              writeOffPtr ptr offset (word8ToAsciiWord8 $ fromIntegral p)
              loop (n - 1) q (offset + 1) True


word8ToAsciiWord8 :: Word8 -> Word8
word8ToAsciiWord8 b = b + fromIntegral (C.ord '0')


word8ToWord64 :: Word8 -> Word64
word8ToWord64 = fromIntegral