{-# LANGUAGE TupleSections #-}
module Codec.AsciiKey256 where

import Control.Applicative
import Control.Monad
import Control.Monad.Fail                     as MF
import Data.Bits
import qualified Data.ByteArray               as BA
         ;import Data.ByteArray               as BA (ByteArrayAccess)
import qualified Data.ByteString              as B
         ;import Data.ByteString              (ByteString)
import qualified Data.ByteString.Base16       as Base16
import qualified Data.ByteString.Base32.Z     as Base32
import qualified Data.ByteString.Base64       as Base64
import qualified Data.ByteString.Char8        as C8
import Data.Char
import Data.Int
import qualified Data.Text                    as T
         ;import Data.Text                    (Text)
import Data.Word
import Foreign.Ptr
import System.IO.Unsafe
import qualified Text.ParserCombinators.ReadP as RP

stripSuffix :: Text -> Text -> Maybe Text
stripSuffix suf x = case T.splitAt (T.length x - T.length suf) x of
    (y,end) | end == suf -> Just y
            | otherwise  -> Nothing

hexdigit :: Char -> Bool
hexdigit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'f') || ( 'A' <= c && c <= 'F')

b64digit :: Char -> Bool
b64digit '.' = True
b64digit '+' = True
b64digit '-' = True
b64digit '/' = True
b64digit c = ('0' <= c && c <= '9') || ( 'a' <= c && c <= 'z') || ( 'A' <= c && c <= 'Z')

-- | Convert to and from a Base64 variant that uses .- instead of +/.
nmtoken64 :: Bool -> Char -> Char
nmtoken64 False '.' = '+'
nmtoken64 False '-' = '/'
nmtoken64 True  '+' = '.'
nmtoken64 True  '/' = '-'
nmtoken64 _     c   = c


-- Apply substitutions for mistaken z-base32 digits.
fixupDigit32 :: Char -> Char
fixupDigit32 'l' = '1'
fixupDigit32 '2' = 'z'
fixupDigit32 'v' = 'u'
fixupDigit32  c  =  c

zb32digit :: Char -> Bool
zb32digit '1' = True
zb32digit c = or [ '3' <= c && c <= '9'
                 , 'a' <= c && c <= 'k'
                 , 'm' <= c && c <= 'u'
                 , 'w' <= c && c <= 'z'
                 ]


-- | Parse 43-digit base64 token into 32-byte bytestring.
parseToken32 :: String -> Either String ByteString
parseToken32 str = fmap (BA.drop 1) $ Base64.decode $ C8.pack $ 'A':map (nmtoken64 False) (take 43 str)

-- | Encode 43-digit base64 token from 32-byte bytestring.
showToken32 :: ByteArrayAccess bin => bin -> String
showToken32 bs = map (nmtoken64 True) $ C8.unpack $ BA.drop 1 $ Base64.encode $ BA.cons 0 $ BA.convert bs

foreign import ccall "crc4itu" c_crc4itu :: Word8        -- ^ init crc
                                            -> Ptr Word8 -- ^ data to checksum
                                            -> Int32     -- ^ length of data
                                            -> IO Word8  -- crc in low 4 bits

-- | CRC4-ITU. Return crc in lowest 4 bits.
crc4itu :: ByteArrayAccess ba => Word8 -- ^ Existing CRC value (usually 0x00) before process a new one.
                                 -> ba -- ^ Data to checksum.
                                 -> Word8
crc4itu crc0 b = unsafePerformIO $ BA.withByteArray b $ \p ->
    c_crc4itu crc0 p (fromIntegral $ BA.length b)

-- | Parse 52-digit z-base32 token into 32-byte bytestring.
parse32Token32 :: String -> Either String ByteString
parse32Token32 str = do
    bs <- Base32.decode $ C8.pack $ map (fixupDigit32 . toLower) (take 52 str) ++ "y"
    case BA.splitAt 32 bs of
        (key,mac) | crc4itu 0 key == shiftR (BA.index mac 0) 4
            -> Right key
        _   -> Left "Failed cyclic redundancy check."

-- | Encode 52-digit z-base32 token from 32-byte bytestring.
show32Token32 :: ByteArrayAccess bin => bin -> String
show32Token32 bs = C8.unpack $ B.take 52 $ Base32.encode (b `B.snoc` shiftL crc 4)
 where
    b = BA.convert bs
    crc = crc4itu 0 bs


readsPrecKey256 :: (ByteString -> Maybe a) -> [Char] -> [(a, [Char])]
readsPrecKey256 publicKey str
        | (bs,_) <- Base16.decode (C8.pack $ take 64 str)
        , Just pub <- publicKey bs
                    = [ (pub, drop (2 * B.length bs) str) ]
        | Right bs <- parse32Token32 str
        , Just pub <- publicKey bs
                    = [ (pub, drop 52 str) ]
        | Right bs <- parseToken32 str
        , Just pub <- publicKey bs
                    = [ (pub, drop 43 str) ]
        | otherwise = []


parseKey256 :: (MonadFail m, Alternative m) => String -> m ByteString
parseKey256 nidstr = do
        let nidbs  = C8.pack nidstr
            (bs,_) = Base16.decode nidbs
            enid   = case C8.length nidbs of
                        52 -> parse32Token32 nidstr
                        43 -> parseToken32 nidstr
                        _  -> Left "Wrong size of key."
        idbs <- (guard (B.length bs == 32) >> return bs)
                <|> either MF.fail return enid
        return idbs

readP_key256 :: RP.ReadP ByteString
readP_key256 = do
      (is64,hexhash) <- foldr1 (RP.+++)
        [ fmap (16,) (sequence $ replicate 64 (RP.satisfy isHexDigit))
        , fmap (32,) (sequence $ replicate 52 (RP.satisfy zb32digit))
        , fmap (64,) (sequence $ replicate 43 (RP.satisfy b64digit))
        ]
      let failure = MF.fail "Bad key."
      case is64 of
        32 -> case parse32Token32 hexhash of
                  Right bs -> return bs
                  _        -> failure
        64 -> case parseToken32 hexhash of
                  Right bs -> return bs
                  _        -> failure
        16 -> case Base16.decode $ C8.pack hexhash of
                  (bs,rem) | B.length bs == 32 &&  B.null rem -> return bs
                  _                                           -> failure
        _ -> failure