MultiBase.hs

-- | Module      : MultiFormats.MultiBase
--   Description : Functions for working with multibase data.
--   Copyright   : Zoey McBride (c) 2026
--   License     : AGPL-3.0-or-later
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
module MultiFormats.MultiBase
  ( MultiBase (prefixChar, baseEncoder, baseDecoder),
    mkMultiBase,

    -- * Decoders
    decodeFromPrefix,
    decodeCIDv0,

    -- * Query functions for MultiBase information
    hasValidPrefix,
    resolvePrefix,
    resolveMultiBase,
  )
where

import Data.BaseSystem (BaseSystem, Decoder, Encoder)
import Data.BaseSystem qualified as Bases
import Data.BaseSystems qualified as Bases
import Data.Function (on)
import Data.List (uncons)
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust)
import Text.Printf (printf)

data MultiBase = MultiBase
  { prefixChar :: Char,
    baseEncoder :: Encoder,
    baseDecoder :: Decoder,
    baseName :: String
  }

instance Eq MultiBase where
  (==) = (==) `on` prefixChar

instance Ord MultiBase where
  compare = compare `on` prefixChar

instance Show MultiBase where
  show mb = printf "multi%s('%c')" (baseName mb) (prefixChar mb)

-- | Implements encoder/decoder for multibase systems when the prefix is known
-- ahead of time.
instance BaseSystem MultiBase where
  encoder mb bytes = prefixChar mb : baseEncoder mb bytes
  decoder _ = decodeFromPrefix

-- | Constructs a Multibase from a BaseSystem.
mkMultiBase :: (Show b, BaseSystem b) => b -> Char -> MultiBase
mkMultiBase bs prefix =
  MultiBase
    { prefixChar = prefix,
      baseEncoder = Bases.encoder bs,
      baseDecoder = Bases.decoder bs,
      baseName = show bs
    }

-- | Resolves a MultiBase system from a Char.
resolvePrefix :: Char -> Maybe MultiBase
resolvePrefix prefix = Map.lookup prefix prefixTable <*> pure prefix
  where
    -- https://github.com/multiformats/multibase/blob/master/multibase.csv
    prefixTable =
      Map.fromList
        [ ('0', mkMultiBase Bases.base2),
          ('9', mkMultiBase Bases.base10),
          ('f', mkMultiBase Bases.base16lower),
          ('F', mkMultiBase Bases.base16upper),
          ('b', mkMultiBase Bases.base32lowerNP),
          ('B', mkMultiBase Bases.base32upperNP),
          ('c', mkMultiBase Bases.base32lower),
          ('C', mkMultiBase Bases.base32upper),
          ('v', mkMultiBase Bases.base32hexlowerNP),
          ('V', mkMultiBase Bases.base32hexupperNP),
          ('t', mkMultiBase Bases.base32hexlower),
          ('T', mkMultiBase Bases.base32hexupper),
          ('z', mkMultiBase Bases.base58btc),
          ('m', mkMultiBase Bases.base64NP),
          ('M', mkMultiBase Bases.base64),
          ('u', mkMultiBase Bases.base64urlNP),
          ('U', mkMultiBase Bases.base64url)
        ]

-- | Resolves a MultiBase system from the first Char in a String.
resolveMultiBase :: String -> Maybe MultiBase
resolveMultiBase str = uncons str >>= \(ch, _) -> resolvePrefix ch

-- | Checks if a given char is an implemented MultiBase prefix.
hasValidPrefix :: String -> Bool
hasValidPrefix [] = False
hasValidPrefix (prefix : _) = isJust $ resolvePrefix prefix

-- | Provides decoder for MultiBase strings with any implemented prefix
-- character.
decodeFromPrefix :: Decoder
decodeFromPrefix str =
  uncons str >>= \(prefix, rest) -> do
    decode <- baseDecoder <$> resolvePrefix prefix
    decode rest

-- | Decodes for CIDv0 from base58btc w/o MultiBase prefix.
decodeCIDv0 :: Decoder
decodeCIDv0 = Bases.decoder Bases.base58btc