{-# LANGUAGE OverloadedStrings #-}
{-|
Module      : Network.Haskoin.Address.Base58
Copyright   : No rights reserved
License     : UNLICENSE
Maintainer  : xenog@protonmail.com
Stability   : experimental
Portability : POSIX

Support for Bitcoin SegWit (BTC) Bech32 addresses. This module is a modified
version of Marko Bencun's reference implementation.
-}
module Network.Haskoin.Address.Bech32
    ( HRP
    , Bech32
    , Data
    , bech32Encode
    , bech32Decode
    , toBase32
    , toBase256
    , segwitEncode
    , segwitDecode
    , Word5(..)
    , word5
    , fromWord5
    ) where

import           Control.Monad         (guard)
import           Data.Array            (Array, assocs, bounds, listArray, (!),
                                        (//))
import           Data.Bits             (Bits, testBit, unsafeShiftL,
                                        unsafeShiftR, xor, (.&.), (.|.))
import qualified Data.ByteString       as B
import           Data.Char             (toUpper)
import           Data.Foldable         (foldl')
import           Data.Functor.Identity (Identity, runIdentity)
import           Data.Ix               (Ix (..))
import           Data.Text             (Text)
import qualified Data.Text             as T
import qualified Data.Text.Encoding    as E
import           Data.Word             (Word8)

-- | Bech32 human-readable string.
type Bech32 = Text

-- | Human-readable part of 'Bech32' address.
type HRP = Text

-- | Data part of 'Bech32' address.
type Data = [Word8]

(.>>.), (.<<.) :: Bits a => a -> Int -> a
(.>>.) = unsafeShiftR
(.<<.) = unsafeShiftL

-- | Five-bit word for Bech32.
newtype Word5 =
    UnsafeWord5 Word8
    deriving (Eq, Ord)

instance Ix Word5 where
    range (UnsafeWord5 m, UnsafeWord5 n) = map UnsafeWord5 $ range (m, n)
    index (UnsafeWord5 m, UnsafeWord5 n) (UnsafeWord5 i) = index (m, n) i
    inRange (m, n) i = m <= i && i <= n

-- | Convert an integer number into a five-bit word.
word5 :: Integral a => a -> Word5
word5 x = UnsafeWord5 (fromIntegral x .&. 31)
{-# INLINE word5 #-}
{-# SPECIALIZE INLINE word5 :: Word8 -> Word5 #-}

-- | Convert a five-bit word into a number.
fromWord5 :: Num a => Word5 -> a
fromWord5 (UnsafeWord5 x) = fromIntegral x
{-# INLINE fromWord5 #-}
{-# SPECIALIZE INLINE fromWord5 :: Word5 -> Word8 #-}

-- | 'Bech32' character map as array of five-bit integers to character.
charset :: Array Word5 Char
charset =
    listArray (UnsafeWord5 0, UnsafeWord5 31) "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

-- | Convert a character to its five-bit value from 'Bech32' 'charset'.
charsetMap :: Char -> Maybe Word5
charsetMap c
    | inRange (bounds inv) upperC = inv ! upperC
    | otherwise = Nothing
  where
    upperC = toUpper c
    inv = listArray ('0', 'Z') (repeat Nothing) // map swap (assocs charset)
    swap (a, b) = (toUpper b, Just a)

-- | Calculate or validate 'Bech32' checksum.
bech32Polymod :: [Word5] -> Word
bech32Polymod values = foldl' go 1 values .&. 0x3fffffff
  where
    go chk value =
        foldl' xor chk' [g | (g, i) <- zip generator [25 ..], testBit chk i]
      where
        generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
        chk' = chk .<<. 5 `xor` fromWord5 value

-- | Convert human-readable part of 'Bech32' string into a list of five-bit
-- words.
bech32HRPExpand :: HRP -> [Word5]
bech32HRPExpand hrp =
    map (UnsafeWord5 . (.>>. 5)) hrpBytes ++
    [UnsafeWord5 0] ++ map word5 hrpBytes
  where
    hrpBytes = B.unpack $ E.encodeUtf8 hrp

-- | Calculate checksum for a string of five-bit words.
bech32CreateChecksum :: HRP -> [Word5] -> [Word5]
bech32CreateChecksum hrp dat = [word5 (polymod .>>. i) | i <- [25,20 .. 0]]
  where
    values = bech32HRPExpand hrp ++ dat
    polymod =
        bech32Polymod (values ++ map UnsafeWord5 [0, 0, 0, 0, 0, 0]) `xor` 1

-- | Verify checksum for a human-readable part and string of five-bit words.
bech32VerifyChecksum :: HRP -> [Word5] -> Bool
bech32VerifyChecksum hrp dat = bech32Polymod (bech32HRPExpand hrp ++ dat) == 1

-- | Maximum length of a Bech32 result.
maxBech32Length :: Int
maxBech32Length = 90

-- | Encode string of five-bit words into 'Bech32' using a provided
-- human-readable part. Can fail if 'HRP' is invalid or result would be longer
-- than 90 characters.
bech32Encode :: HRP -> [Word5] -> Maybe Bech32
bech32Encode hrp dat = do
    guard $ checkHRP hrp
    let dat' = dat ++ bech32CreateChecksum hrp dat
        rest = map (charset !) dat'
        result = T.concat [T.toLower hrp, T.pack "1", T.pack rest]
    guard $ T.length result <= maxBech32Length
    return result

-- | Check that human-readable part is valid for a 'Bech32' string.
checkHRP :: HRP -> Bool
checkHRP hrp = not (T.null hrp) && T.all (\char -> char >= '\x21' && char <= '\x7e') hrp

-- | Decode human-readable 'Bech32' string into a human-readable part and a
-- string of five-bit words.
bech32Decode :: Bech32 -> Maybe (HRP, [Word5])
bech32Decode bech32 = do
    guard $ T.length bech32 <= maxBech32Length
    guard $ T.toUpper bech32 == bech32 || lowerBech32 == bech32
    let (hrp, dat) = T.breakOnEnd "1" lowerBech32
    guard $ T.length dat >= 6
    hrp' <- T.stripSuffix "1" hrp
    guard $ checkHRP hrp'
    dat' <- mapM charsetMap $ T.unpack dat
    guard $ bech32VerifyChecksum hrp' dat'
    return (hrp', take (T.length dat - 6) dat')
  where
    lowerBech32 = T.toLower bech32

type Pad f = Int -> Int -> Word -> [[Word]] -> f [[Word]]

yesPadding :: Pad Identity
yesPadding _ 0 _ result        = return result
yesPadding _ _ padValue result = return $ [padValue] : result
{-# INLINE yesPadding #-}

noPadding :: Pad Maybe
noPadding frombits bits padValue result = do
    guard $ bits < frombits && padValue == 0
    return result
{-# INLINE noPadding #-}

-- | Big endian conversion of a bytestring from base \(2^{frombits}\) to base
-- \(2^{tobits}\). {frombits} and {twobits} must be positive and
-- \(2^{frombits}\) and \(2^{tobits}\) must be smaller than the size of Word.
-- Every value in 'dat' must be strictly smaller than \(2^{frombits}\).
convertBits :: Functor f => [Word] -> Int -> Int -> Pad f -> f [Word]
convertBits dat frombits tobits pad = concat . reverse <$> go dat 0 0 []
  where
    go [] acc bits result =
        let padValue = (acc .<<. (tobits - bits)) .&. maxv
        in pad frombits bits padValue result
    go (value:dat') acc bits result =
        go dat' acc' (bits' `rem` tobits) (result' : result)
      where
        acc' = (acc .<<. frombits) .|. fromIntegral value
        bits' = bits + frombits
        result' =
            [ (acc' .>>. b) .&. maxv
            | b <- [bits' - tobits,bits' - 2 * tobits .. 0]
            ]
    maxv = (1 .<<. tobits) - 1
{-# INLINE convertBits #-}

-- | Convert from eight-bit to five-bit word string, adding padding as required.
toBase32 :: [Word8] -> [Word5]
toBase32 dat =
    map word5 $ runIdentity $ convertBits (map fromIntegral dat) 8 5 yesPadding

-- | Convert from five-bit word string to eight-bit word string, ignoring padding.
toBase256 :: [Word5] -> Maybe [Word8]
toBase256 dat =
    map fromIntegral <$> convertBits (map fromWord5 dat) 5 8 noPadding

-- | Check if witness version and program are valid.
segwitCheck :: Word8 -> Data -> Bool
segwitCheck witver witprog =
    witver <= 16 &&
    if witver == 0
        then length witprog == 20 || length witprog == 32
        else length witprog >= 2 && length witprog <= 40

-- | Decode SegWit 'Bech32' address from a string and expected human-readable part.
segwitDecode :: HRP -> Bech32 -> Maybe (Word8, Data)
segwitDecode hrp addr = do
    (hrp', dat) <- bech32Decode addr
    guard $ (hrp == hrp') && not (null dat)
    let (UnsafeWord5 witver:datBase32) = dat
    decoded <- toBase256 datBase32
    guard $ segwitCheck witver decoded
    return (witver, decoded)

-- | Encode 'Data' as a SegWit 'Bech32' address. Needs human-readable part and
-- witness program version.
segwitEncode :: HRP -> Word8 -> Data -> Maybe Text
segwitEncode hrp witver witprog = do
    guard $ segwitCheck witver witprog
    bech32Encode hrp $ UnsafeWord5 witver : toBase32 witprog