Supplemental.hs

-- | Module      : IPLD.DagCBOR.Decoder.Supplemental
--   Description : Parsing additional data from argument types.
--   Copyright   : Zoey McBride (c) 2026
--   License     : AGPL-3.0-or-later
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
module IPLD.DagCBOR.Decoder.Supplemental
  ( -- * Type-class for implementing Supplements for Arguments.
    Supplemental (decodeSup),

    -- * Implements CBOR simple values (Bools, Nulls, FloatData)
    SimpleData (..),

    -- * Types for working with CBOR binary objects
    ListObj,
    MapObj,

    -- * DecoderMethods for supplementals
    supplementIntData,
  )
where

import Data.Bifunctor (Bifunctor (bimap, first))
import Data.Bits ((.<<.), (.|.))
import Data.ByteString qualified as LBS
import Data.ByteString.Builder qualified as Builder
import Data.Text (Text)
import Data.Text.Encoding qualified as Text
import Data.XCodec.BinaryTranscoder (BinaryTranscoder)
import Data.XCodec.BinaryTranscoder qualified as BXC
import IPLD.DM qualified as DM
import IPLD.DagCBOR.Decoder.Arguments
import IPLD.DagCBOR.Decoder.Errors
import IPLD.DagCBOR.Decoder.MajorTypes
import MultiFormats.CID (CID)
import MultiFormats.CID.Constants qualified as CID
import MultiFormats.CID.Extractor qualified as CID
import Unsafe.Coerce (unsafeCoerce)

-- | Defines arguments which have additional data after the parameter.
class (Argument a, BinaryTranscoder b) => Supplemental a t b where
  -- | Provides function for reading argument data from argument type.
  decodeSup :: a -> Okay (DecoderMethod t b)

-- | Gives back the bytes converted into `a` when minimally encoded.
supplementIntData :: (BinaryTranscoder bxc) => Int -> DecoderMethod IntData bxc
supplementIntData sizeof bytes = do
  -- Load the #bytes for the integer
  Result intdata rest <- extractBytes sizeof bytes
  -- Build the data from big-endian encoded CBOR
  let value = fromIntegral $ BXC.unpackValueBE intdata
      result = Result value rest
  -- Match the size with its minimal encoding check
  case sizeof of
    1 -- Check 8-bit encodings
      | not (minimal8 value) -> Left NotMinimal
      | otherwise -> Right result
    2 -- Check 16-bit encodings
      | not (minimal16 value) -> Left NotMinimal
      | otherwise -> Right result
    4 -- Check 32-bit encodings
      | not (minimal32 value) -> Left NotMinimal
      | otherwise -> Right result
    8 -- Check 64-bit minimal encoding
      | not (minimal64 value) -> Left NotMinimal
      | otherwise -> Right result
    _ -> error "supplementIntData: sizeof invariant failed"
  where
    -- Valid minimal-encoding ranges for different byte-sized integers
    minimal8 r = 0 <= r && r <= 0xFF
    minimal16 r = 0xFF < r && r <= 0xFFFF
    minimal32 r = 0xFFFF < r && r <= 0xFFFFFFFF
    minimal64 r = 0xFFFFFFFF < r && r <= 0xFFFFFFFFFFFFFFFF

-- | Implements reading a Number after an Ints argument. All numbers get stored
-- in a 64-bit integer for performance.
instance (BinaryTranscoder bxc) => Supplemental Ints IntData bxc where
  decodeSup arg =
    case arg of
      Ints5Bits value -> Right $ paramIntData value
      Ints8Bits -> Right $ supplementIntData 1
      Ints16Bits -> Right $ supplementIntData 2
      Ints32Bits -> Right $ supplementIntData 4
      Ints64Bits -> Right $ supplementIntData 8
    where
      -- Creates a Result from a 5-bit int parameter.
      paramIntData value bytes =
        Right $ Result (fromIntegral value :: IntData) bytes

-- | Implements reading a Supplemental collection of binary transcoder data.
instance (BinaryTranscoder b) => Supplemental Collection b b where
  decodeSup arg =
    case arg of
      CollectionsIndefinite -> Left Indefinite
      CollectionsLength ints ->
        Right $ \bytes -> do
          -- Get the function for reading the int
          intdata <- decodeSup ints :: Okay (DecoderMethod IntData b)
          -- Read the #bytes from the IntData
          Result count rest <- intdata bytes
          extractBytes (fromIntegral count) rest

-- | Implements reading a Supplemental CBOR bytestring data.
instance (BinaryTranscoder b) => Supplemental Collection DM.BytesData b where
  decodeSup arg =
    Right $ \bytes -> do
      -- Read the bytes from the transcoder data
      readbytes <- decodeSup arg :: Okay (DecoderMethod b b)
      Result bytesdata rest <- readbytes bytes
      -- Serialize the binary value into Builder
      let builder = BXC.unpackSerial bytesdata
          len = BXC.totalOctets bytesdata
      Right $ Result (DM.BytesData builder len) rest

-- | Implements reading a Supplemental collection of UTF-8 text.
instance (BinaryTranscoder bxc) => Supplemental Collection Text bxc where
  decodeSup arg =
    Right $ \bytes -> do
      -- Get the method for reading ByteStrings
      readbytes <- decodeSup arg :: Okay (DecoderMethod DM.BytesData bxc)
      -- Read the decoder result
      Result bytesdata rest <- readbytes bytes
      -- Decode the result
      Right $ Result (decodeUtf8 bytesdata) rest
    where
      -- Instansiates the Builder into a strict ByteString for Text.decodeUtf8
      decodeUtf8 (DM.BytesData builder _) =
        Text.decodeUtf8 $ LBS.toStrict (Builder.toLazyByteString builder)

-- | Iterates some monadic action `f` on the pure `a` value at most `n` times.
iterateTakeM :: (Monad m) => Int -> (a -> m a) -> a -> m [a]
iterateTakeM n f = sequence . take n . drop 1 . iterate (f =<<) . pure

-- | Stores length information for a parsed CBOR object.
data Obj bxc = Obj IntData bxc

-- | Removes the CBOR object from the bytes.
nextObj :: (BinaryTranscoder bxc) => bxc -> Okay (Obj bxc, bxc)
nextObj bytes = do
  Result size _ <- majorTypeSize bytes
  isize <- downcastIntData size
  let (hi, lo) = BXC.splitOffset isize bytes
  Right (Obj size hi, lo)

-- | Folds the resulting list of CBOR objects from the iteration list.
foldObjs :: bxc -> [(Obj a, bxc)] -> Result (ListObj a) bxc
foldObjs start = foldr buildResult (Result (0, []) start)
  where
    -- Returns true if the result is the initial accumulator to foldr
    isInitial (Result (size, _) _) = size == 0
    -- Builds the Result from the (Obj a, bxc) list fold.
    buildResult (Obj len bxc, afterobjs) res
      -- If this is the first item, this is the bytestring with all the objects
      -- removed, so update the pointer in the second item in Result.
      | isInitial res = bimap (bimap (len +) (bxc :)) (const afterobjs) res
      -- Otherwise, only update the length and list of BXCs
      | otherwise = first (bimap (len +) (bxc :)) res

-- | Stores a list of CBOR binary objects + length.
type ListObj bxc = (IntData, [bxc])

-- | Chunks DAG-CBOR Lists into list of encoded binary CBOR objects.
instance (BinaryTranscoder b) => Supplemental Collection (ListObj b) b where
  decodeSup arg =
    case arg of
      CollectionsIndefinite -> Left Indefinite
      CollectionsLength ints ->
        Right $ \bytes -> do
          -- Read the list length data
          Result elemcount start <- extractIntData ints bytes
          ielems <- downcastIntData elemcount
          -- Iterate through CBOR list members for the #elements
          foldObjs start <$> iterateTakeM ielems nextElem (Obj 0 mempty, start)
    where
      -- Extracts the next element from the list with iterate splitAt
      nextElem (Obj total _, bytes) = do
        (Obj len obj, rest) <- nextObj bytes
        return (Obj (len + total) obj, rest)

-- | Stores key-value CBOR binary object pairs + length.
type MapObj bxc = ListObj (bxc, bxc)

-- | Chunks DAG-CBOR Map into list of encoded binary CBOR key-value pairs.
instance (BinaryTranscoder b) => Supplemental Collection (MapObj b) b where
  decodeSup arg =
    case arg of
      CollectionsIndefinite -> Left Indefinite
      CollectionsLength ints ->
        Right $ \bytes -> do
          -- Read the number of pairs for the map data
          Result len start <- extractIntData ints bytes
          ipairs <- downcastIntData len
          -- Iterate through CBOR list members for the #pairs
          foldObjs start <$> iterateTakeM ipairs nextPair (mkIterObj start)
    where
      -- Initializes the iterator for collecting key value pairs
      mkIterObj start = (Obj 0 (mempty, mempty), start)
      -- Extracts the next pair from the map with iterate splitAt
      nextPair (Obj total _, bytes) = do
        (Obj l1 key, onval) <- nextObj bytes
        (Obj l2 val, rest) <- nextObj onval
        return (Obj (total + l1 + l2) (key, val), rest)

-- | Implements reading a Supplemental CID from the matching tag, see:
-- https://github.com/ipld/cid-cbor/
-- https://ipld.io/specs/codecs/dag-cbor/spec/#links
instance (BinaryTranscoder bxc) => Supplemental Tag CID bxc where
  decodeSup Tag =
    Right $ \bytes -> do
      -- Remove the CID tag identifier
      Result _ inner1 <- removeOctet CID.ipldTagCBOR bytes
      -- Decode the arg from the first byte and get the decoder rule
      Result arg inner2 <- extractByte inner1
      -- Get the rule for reading in the CID data bytestring.
      bytesrule <- decodeArg arg :: Okay Collection
      bytesread <- decodeSup bytesrule :: Okay (DecoderMethod bxc bxc)
      -- Remove and verify the multibase 0 prefix and extract the CID result.
      Result cidtop final <- bytesread inner2
      Result _ ciddata <- removeOctet 0 cidtop
      case CID.extractor ciddata of
        Left err -> Left $ CIDError err
        Right (cid, fragmented)
          | BXC.totalOctets fragmented > 0 -> Left Fragmented
          | otherwise -> Right $ Result cid final

-- | Reads simple values. For floating points, it does:
-- 1. Takes n number bytes from the ByteString
-- 2. Validates bytes taken to matches the FloatData type size.
-- 3. Copy those bytes to a corresponding sized Word type
-- 4. Check if they are a valid IPLD Float (Not -∞, ∞, or NaN)
-- 5. Updates the bytes to remove the floating point bytes.
-- Otherwise, it packages the simple value and leaves the bytes unmodified.
instance (BinaryTranscoder bxc) => Supplemental Simple SimpleData bxc where
  decodeSup arg =
    case arg of
      -- Floating point values
      FloatHalf -> Right $ float DM.HalfPrecision 2
      FloatSingle -> Right $ float DM.SinglePrecision 4
      FloatDouble -> Right $ float DM.DoublePrecision 8
      -- Simple supported values
      SimpleNull -> Right $ simple NullData
      SimpleBool bool -> Right $ simple (BoolData bool)
      -- Simple unsupported values
      SimpleUndefined -> Left NotDefined
      StopCollection -> Left Indefinite
    where
      -- Returns the simple value without changing the bytes.
      simple value bytes = Right $ Result value bytes
      -- Copys and constructs a floating point number from n bytes.
      -- TODO: I want to be able to relax the validation with an ifdef
      float mk n bytes = do
        -- Take the first n bytes from the ByteString
        Result taken rest <- extractBytes n bytes
        -- Convert the value from network order to CPU's format.
        value <- copyBEtoCPU taken
        -- Validate its IEEE754 representation. Use `unsafeCoerce` to change
        -- the bit interpretation from Word to Float/Double.
        case (unsafeCoerce value, unsafeCoerce value) of
          (as32, as64)
            -- NOTE: 16-bit floating point numbers aren't part of Haskell.
            | n == 2 -> Left $ NotImplemented "decodeSup: 16-bit float"
            -- Validate 32-bit single floats
            | n == 4 && invalidSingle as32 -> Left InvalidIEEE754
            -- Validate 64-bit double floats
            | n == 8 && invalidDouble as64 -> Left InvalidIEEE754
            -- Success: return the constructed value
            | otherwise -> Right $ Result (FloatData $ mk value) rest
        where
          -- Copys from BE format to CPU format
          -- WARN: only supports little-endian CPUs.
          -- TODO: there should be a way to convert endianness in XCodec
          copyBEtoCPU = copyW8s 0 0 . BXC.unpackOctets . BXC.swapOrder
          -- Bitwise-or each byte from the number into the Word type. Fail when
          -- iteration exceeds bytes or doesn't read the correct #bytes.
          copyW8s !i !acc (!w8 : ws)
            | i < n = copyW8s (i + 1) (acc .<<. 8 .|. fromIntegral w8) ws
            | otherwise = Left UnexpectedStop
          copyW8s !i !acc []
            | i == n = Right acc
            | otherwise = Left UnexpectedStop
          -- True if invalid 32-bit value
          invalidSingle (single :: Float) = isNaN single || isInfinite single
          -- True if invalid 64-bit value
          invalidDouble (double :: Double) = isNaN double || isInfinite double