{-# LANGUAGE StrictData #-}
module Bitcoin.Address.SegWit
  ( -- * Version
    Version
  , version
  , unVersion
  , versionOp
  , version0
    -- * Programs
  , Program
  , program
  , programVersion
  , programData
  , renderProgram
    -- ** Standard programs
  , p2wpkh
  , p2wsh
  ) where

import qualified Codec.Binary.Bech32 as Bech32
import qualified Codec.Binary.Bech32.Internal as Bech32
import Control.Monad
import qualified Data.Bitcoin.Script as S
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy.Char8 as BL8
import qualified Data.Text.Encoding as T
import Data.Word

import Bitcoin.Address.Hash
import Bitcoin.Address.Internal (op0to16)
import Bitcoin.Address.Settings

--------------------------------------------------------------------------------

-- | A SegWit program. Construct with 'program'.
data Program = Program Version B.ByteString
  deriving (Eq, Ord)

-- | Version and base-16 encoded program data.
instance Show Program where
  showsPrec n (Program v d) = showParen (n > 10) $
    showString "Program " .
    mappend (show v) .
    mappend " " .
    mappend (BL8.unpack (BB.toLazyByteString (BB.byteStringHex d)))

-- | SegWit program version.
programVersion :: Program -> Version
{-# INLINE programVersion #-}
programVersion (Program v _) = v

-- | Raw SegWit program data.
programData :: Program -> B.ByteString
{-# INLINE programData #-}
programData (Program _ x) = x

-- | Construct a 'Program' from its raw bytes.
--
-- __WARINING__ This function will prevent you frow constructing invalid SegWit
-- programs, but won't help you write __meaningful__ programs. Prefer to use
-- safe constructions such as 'p2wpkh' or 'p2wsh' instead.
program
  :: Version
  -> B.ByteString  -- ^ Raw SegWit program bytes.
  -> Maybe Program -- ^ Nothing if program length is invalid for version.
program ver prog = do
  let len = B.length prog
  guard $ case unVersion ver of
     0 -> len == 20 || len == 32
     _ -> len >= 2  && len <= 40
  pure (Program ver prog)

--------------------------------------------------------------------------------

-- | Bech32-encode a 'Program'.
renderProgram :: PrefixSegWit -> Program -> B.ByteString
{-# INLINE renderProgram #-}
renderProgram pre (Program ver prog) =
  let w5ver = toEnum $ fromIntegral $ unVersion ver
      dver  = Bech32.dataPartFromWords [w5ver]
      dprog = Bech32.dataPartFromBytes prog
  in T.encodeUtf8 $ Bech32.encodeLenient (prefixSegWitHRP pre) (dver <> dprog)

--------------------------------------------------------------------------------

-- | The version for of a SegWit 'Program'.
newtype Version = Version Word8
  deriving (Eq, Ord, Show)

-- | Construct a SegWit 'Version'.
--
-- The given 'Word8' must be in the inclusive range [0 … 16].
version :: Word8 -> Maybe Version
{-# INLINE version #-}
version w = do
  guard (w <= 16)
  Just (Version w)

-- | The obtained 'Word8' is in the inclusive range [0 … 16].
unVersion :: Version -> Word8
{-# INLINE unVersion #-}
unVersion (Version w) = w

-- | The 'S.ScriptOp' corresponding to the 'Version', in
-- range ['S.OP_0' … 'S.OP_16']
versionOp :: Version -> S.ScriptOp
{-# INLINE versionOp #-}
versionOp (Version w) = case op0to16 (fromIntegral w) of
  Just op -> op
  Nothing -> undefined -- impossible

--------------------------------------------------------------------------------

-- | SegWit version 0
version0 :: Version
{-# INLINE version0 #-}
Just version0 = version 0

-- | Construct a standard SegWit version 0 P2WPKH program.
p2wpkh :: PubHash160 -> Program
{-# INLINE p2wpkh #-}
p2wpkh pkh = case program version0 (unPubHash160 pkh) of
  Just prog -> prog
  Nothing   -> undefined -- impossible

-- | Construct a standard SegWit version 0 P2WSH program.
p2wsh :: ScriptSHA256 -> Program
{-# INLINE p2wsh #-}
p2wsh sh = case program version0 (unScriptSHA256 sh) of
  Just prog -> prog
  Nothing   -> undefined -- impossible