CID.hs

-- | Module      : MultiFormats.CID
--   Description : Implements multiformats CID
--   Copyright   : Zoey McBride (c) 2026
--   License     : AGPL-3.0-or-later
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
--
-- This module provides a way for extracting, parsing and serializing CID data.
-- See: https://github.com/multiformats/cid
module MultiFormats.CID
  ( CID (CID),
    cidVersion,
    cidDataCodec,
    cidHashCodec,
    cidHashBuilder,
    cidHashSize,
    cidSizeOf,
  )
where

import Control.DeepSeq (NFData (..), force, ($!!))
import Data.BaseSystems.Lazy (base16lower, encoder)
import Data.Bifunctor (first)
import Data.Bits ((.>>.))
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Short (ShortByteString)
import Data.List (intercalate)
import Data.Maybe (fromJust)
import Data.Text qualified as Text
import Data.XCodec.BinaryTranscoder (BinaryTranscoder)
import Data.XCodec.BinaryTranscoder qualified as BXC
import MultiFormats.CID.Component.Codec qualified as CID
import MultiFormats.CID.Component.Decoder qualified as CID
import MultiFormats.CID.Component.Hash qualified as CID
import MultiFormats.CID.Component.Version (Version (CIDv0, CIDv1))
import MultiFormats.CID.Component.Version qualified as CID
import MultiFormats.CID.Constants qualified as CID
import MultiFormats.CID.Errors
import MultiFormats.CID.Extractor
import MultiFormats.CID.Parser
import MultiFormats.CID.Serializer
import MultiFormats.MultiCodec (MultiCodec (..), resolveMultiCodec)
import MultiFormats.MultiHash (MultiHash (..), resolveMultiHash)
import MultiFormats.VarInt (VarInt)

-- | Defines the data-type for CID data.
data CID = CID
  { -- | Specifies the version (CIDv0 or CIDv1) of this CID.
    cidVersion :: CID.Version,
    -- | Describes the content type of the hashed data.
    cidDataCodec :: MultiCodec,
    -- | Contains the hash function used.
    cidHashCodec :: MultiHash,
    -- | Stores the # of bytes used to store the HashDigest
    cidHashSize :: VarInt,
    -- | Stores the hash function's output.
    cidHashBuilder :: Builder
  }

-- | Returns the size of a CID in bytes.
{-# INLINE cidSizeOf #-}
cidSizeOf :: CID -> VarInt
cidSizeOf (CID CIDv0 _ _ _ _) = fromIntegral CID.v0TotalBytes
cidSizeOf (CID CIDv1 datacodec hashcodec hashsize _) =
  let datavar = multiCodecValue datacodec
      hashvar = multiCodecValue $ multiHashCodec hashcodec
   in hashsize
        + varSizeOf 0 CID.v1FirstVarInt
        + varSizeOf 0 datavar
        + varSizeOf 0 hashsize
        + varSizeOf 0 hashvar
  where
    -- Gets the #bytes for a VarInt, including the extra byte for 8-byte values.
    varSizeOf !total !var
      -- Shift the varint by 1 varint group (7 bits) until it reaches zero
      | var /= 0 = varSizeOf (total + 1) (var .>>. 7)
      -- Adjust the total for when the byte value is zero.
      | otherwise = max 1 total

-- | Checks equality of the CID values, version, and hash digest.
instance Eq CID where
  (CID v1 c1 h1 s1 d1) == (CID v2 c2 h2 s2 d2) =
    -- Check that none of the comparable values mismatch
    not (v1 /= v2 || c1 /= c2 || h1 /= h2 || s1 /= s2)
      -- And then check equality on the executed builder
      && Builder.toLazyByteString d1 == Builder.toLazyByteString d2

-- | Implements Show for CID.
instance Show CID where
  show (CID version datacodec hashcodec hashsize hashbuilder) =
    let hashdigest = Builder.toLazyByteString hashbuilder
        hashstring = "0x" ++ Text.unpack (encoder showbase hashdigest)
     in show version
          ++ "("
          ++ intercalate
            ", "
            [ "DataCodec=" ++ show datacodec,
              "HashCodec=" ++ show hashcodec,
              "HashDigest[" ++ show hashsize ++ "]=" ++ hashstring
            ]
          ++ ")"
    where
      -- Sets the base implementation to show the hash digest with.
      showbase = base16lower

-- | Evaluates CID struct into normal form without executing the hash builder.
instance NFData CID where
  rnf (CID version datacodec hashcodec hashsize _) =
    rnf version `seq` rnf datacodec `seq` rnf hashcodec `seq` rnf hashsize

-- | Parses a CIDv1 from a MultiBase encoded string OR a CIDv0 from a base16lower
-- encoded string. Fails when there is extra bytes past the CID.
instance Parser CID where
  parser cidstr = do
    -- Decode to bytes based on the String's encoding. Use SBS as BXC because
    -- the bytestring is discarded after this function.
    CID.Decoder cidbytes <- parser cidstr :: Okay (CID.Decoder ShortByteString)
    -- Extracts the CID's data from bytes.
    (cid, trailing) <- extractor cidbytes
    -- Fail if there is any leftovers. Otherwise fully evaluate the CID.
    if BXC.totalOctets trailing > 0
      then Left (InvalidData TrailingBytes)
      else Right $!! cid

-- | Extracts a CID from a ByteString. Allows trailing bytes.
instance (BinaryTranscoder t) => Extractor CID t where
  extractor bytes = do
    -- First, extract the version from bytes
    (version, rest) <- extractor bytes
    -- Apply the constructor if the CID was successfully extracted from bytes,
    -- and then strictly evaluate the CID into memory.
    first force <$> case version of
      CIDv0 -> extractCIDv0 rest
      CIDv1 -> extractCIDv1 rest

-- | Applies the extraction after version# for CIDv0 was extracted.
extractCIDv0 :: (BinaryTranscoder t) => Extraction CID t
extractCIDv0 ondata
  -- Fail if there isn't enough bytes
  | totalbytes /= CID.v0HashBytes = Left (InvalidHash IncorrectLength)
  -- Create the CIDv0 from data
  | otherwise = Right (mkCIDv0 ciddata, rest)
  where
    totalbytes = BXC.totalOctets ondata
    (ciddata, rest) = BXC.splitOffset CID.v0HashBytes ondata

-- | Applies the extraction after version# for CIDv1 was extracted.
extractCIDv1 :: (BinaryTranscoder t) => Extraction CID t
extractCIDv1 oncodec = do
  -- Extract the 1st MultiCodec from the bytes.
  (CID.Codec multicodec, onhash) <- extractor oncodec
  -- Extract the 2nd MultiCodec and get CIDHash
  (CID.Hash multihash hashsize hashdigest, stop) <- extractor onhash
  -- Get the length of the hash data after the hash codec
  -- TODO: verifiy hash length against crypton library
  return (CID CIDv1 multicodec multihash hashsize hashdigest, stop)

-- | Serializes a CID from representation.
instance Serializer CID where
  serializer ciddata =
    case cidVersion ciddata of
      CIDv0 -> serializeCIDv0 ciddata
      CIDv1 -> serializeCIDv1 ciddata

-- | Serializes CIDv0 by adding the Codec Byte to its digest
serializeCIDv0 :: Serialization CID
serializeCIDv0 (CID CIDv1 _ _ _ _) = error "serializeCIDv0: input has CIDv1!"
serializeCIDv0 (CID CIDv0 _ _ _ cidhash) =
  -- Concat the Builders for the bytes data into a single Builder
  let cid1st = Builder.word8 CID.v0FirstByte
      cid2nd = Builder.word8 CID.v0SecondByte
   in return $ mconcat [cid1st, cid2nd, cidhash]

-- | Serializes CIDv1 into VarInts of its components.
serializeCIDv1 :: Serialization CID
serializeCIDv1 (CID CIDv0 _ _ _ _) = error "serializeCIDv1: input has CIDv0!"
serializeCIDv1 (CID CIDv1 datacodec hashcodec hashlength hashbuilder) = do
  -- Get the Builder from each serializer sequenced from Just values
  builders <-
    sequence
      [ serializer CID.v1FirstVarInt,
        serializer datacodec,
        serializer hashcodec,
        serializer hashlength,
        pure hashbuilder
      ]
  -- Concat the list into one large Builder
  return $ mconcat builders

-- | Provides the contructor for CIDv0s, see:
-- https://github.com/multiformats/cid?tab=readme-ov-file#cidv0
mkCIDv0 :: (BinaryTranscoder t) => t -> CID
mkCIDv0 digest =
  CID
    { cidVersion = CIDv0,
      cidDataCodec = codec "dag-pb",
      cidHashCodec = hash "sha2-256",
      cidHashSize = varlen digest,
      cidHashBuilder = BXC.unpackSerial digest
    }
  where
    codec str = fromJust $!! resolveMultiCodec str
    hash str = fromJust $!! resolveMultiHash str
    varlen bxc = fromIntegral $! BXC.totalOctets bxc