{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}
{-|
Module      : Haskoin.Address
Copyright   : No rights reserved
License     : MIT
Maintainer  : jprupp@protonmail.ch
Stability   : experimental
Portability : POSIX

Base58, CashAddr, Bech32 address and WIF private key serialization support.
-}
module Haskoin.Address
    ( -- * Addresses
      Address(..)
    , isPubKeyAddress
    , isScriptAddress
    , isWitnessPubKeyAddress
    , isWitnessScriptAddress
    , addrToText
    , textToAddr
    , bech32ToAddr
    , cashToAddr
    , base58ToAddr
    , addrToJSON
    , addrToEncoding
    , addrFromJSON
    , pubKeyAddr
    , pubKeyWitnessAddr
    , pubKeyCompatWitnessAddr
    , p2pkhAddr
    , p2wpkhAddr
    , p2shAddr
    , p2wshAddr
    , inputAddress
    , outputAddress
    , addressToScript
    , addressToScriptBS
    , addressToOutput
    , payToScriptAddress
    , payToWitnessScriptAddress
    , payToNestedScriptAddress
    , scriptToAddress
    , scriptToAddressBS
    , module Haskoin.Address.Base58
    , module Haskoin.Address.Bech32
    , module Haskoin.Address.CashAddr
    ) where

import           Control.Applicative
import           Control.Arrow            (second)
import           Control.DeepSeq
import           Control.Monad
import           Data.Aeson               as A
import           Data.Aeson.Encoding      as A
import           Data.Aeson.Types
import           Data.Binary              (Binary (..))
import           Data.ByteString          (ByteString)
import qualified Data.ByteString          as B
import           Data.Bytes.Get
import           Data.Bytes.Put
import           Data.Bytes.Serial
import           Data.Hashable
import           Data.Maybe
import           Data.Serialize           (Serialize (..))
import           Data.Text                (Text)
import qualified Data.Text                as T
import           Data.Word                (Word8)
import           GHC.Generics             (Generic)
import           Haskoin.Address.Base58
import           Haskoin.Address.Bech32
import           Haskoin.Address.CashAddr
import           Haskoin.Constants
import           Haskoin.Crypto
import           Haskoin.Keys.Common
import           Haskoin.Script
import           Haskoin.Util

-- | Address format for Bitcoin and Bitcoin Cash.
data Address
    -- | pay to public key hash (regular)
    = PubKeyAddress
          { getAddrHash160 :: !Hash160
          -- ^ RIPEMD160 hash of public key's SHA256 hash
          }
    -- | pay to script hash
    | ScriptAddress
          { getAddrHash160 :: !Hash160
          -- ^ RIPEMD160 hash of script's SHA256 hash
          }
    -- | pay to witness public key hash
    | WitnessPubKeyAddress
          { getAddrHash160 :: !Hash160
          -- ^ RIPEMD160 hash of public key's SHA256 hash
          }
    -- | pay to witness script hash
    | WitnessScriptAddress
          { getAddrHash256 :: !Hash256
          -- ^ HASH256 hash of script
          }
    -- | other witness address
    | WitnessAddress
         { getAddrVersion :: !Word8
         , getAddrData    :: !ByteString
         }
    deriving
        (Eq, Ord, Generic, Show, Read, Hashable, NFData)

instance Serial Address where
    serialize (PubKeyAddress k) = do
        putWord8 0x00
        serialize k
    serialize (ScriptAddress s) = do
        putWord8 0x01
        serialize s
    serialize (WitnessPubKeyAddress h) = do
        putWord8 0x02
        serialize h
    serialize (WitnessScriptAddress s) = do
        putWord8 0x03
        serialize s
    serialize (WitnessAddress v d) = do
        putWord8 0x04
        putWord8 v
        putWord64be (fromIntegral (B.length d))
        putByteString d

    deserialize =
        getWord8 >>= \case
        0x00 -> PubKeyAddress <$> deserialize
        0x01 -> ScriptAddress <$> deserialize
        0x02 -> WitnessPubKeyAddress <$> deserialize
        0x03 -> WitnessScriptAddress <$> deserialize
        0x04 -> WitnessAddress <$> getWord8 <*>
                (getByteString . fromIntegral =<< getWord64be)
        b -> fail . T.unpack $
             "Could not decode address type byte: " <>
             encodeHex (B.singleton b)

instance Serialize Address where
    put = serialize
    get = deserialize

instance Binary Address where
    put = serialize
    get = deserialize

-- | 'Address' pays to a public key hash.
isPubKeyAddress :: Address -> Bool
isPubKeyAddress PubKeyAddress {} = True
isPubKeyAddress _                = False

-- | 'Address' pays to a script hash.
isScriptAddress :: Address -> Bool
isScriptAddress ScriptAddress {} = True
isScriptAddress _                = False

-- | 'Address' pays to a witness public key hash. Only valid for SegWit
-- networks.
isWitnessPubKeyAddress :: Address -> Bool
isWitnessPubKeyAddress WitnessPubKeyAddress {} = True
isWitnessPubKeyAddress _                       = False

isWitnessScriptAddress :: Address -> Bool
isWitnessScriptAddress WitnessScriptAddress {} = True
isWitnessScriptAddress _                       = False

isWitnessAddress :: Address -> Bool
isWitnessAddress WitnessAddress {} = True
isWitnessAddress _                 = False

addrToJSON :: Network -> Address -> Value
addrToJSON net a = toJSON (addrToText net a)

addrToEncoding :: Network -> Address -> Encoding
addrToEncoding net = maybe null_ text . addrToText net

-- | JSON parsing for Bitcoin addresses. Works with 'Base58', 'CashAddr' and
-- 'Bech32'.
addrFromJSON :: Network -> Value -> Parser Address
addrFromJSON net =
    withText "address" $ \t ->
        case textToAddr net t of
            Nothing -> fail "could not decode address"
            Just x  -> return x

-- | Convert address to human-readable string. Uses 'Base58', 'Bech32', or
-- 'CashAddr' depending on network.
addrToText :: Network -> Address -> Maybe Text
addrToText net a@PubKeyAddress {getAddrHash160 = h}
    | isNothing (getCashAddrPrefix net) =
        Just . encodeBase58Check . runPutS $ base58put net a
    | otherwise = cashAddrEncode net 0 (runPutS $ serialize h)
addrToText net a@ScriptAddress {getAddrHash160 = h}
    | isNothing (getCashAddrPrefix net) =
        Just . encodeBase58Check . runPutS $ base58put net a
    | otherwise =
        cashAddrEncode net 1 (runPutS $ serialize h)
addrToText net WitnessPubKeyAddress {getAddrHash160 = h} = do
    hrp <- getBech32Prefix net
    segwitEncode hrp 0 (B.unpack (runPutS $ serialize h))
addrToText net WitnessScriptAddress {getAddrHash256 = h} = do
    hrp <- getBech32Prefix net
    segwitEncode hrp 0 (B.unpack (runPutS $ serialize h))
addrToText net WitnessAddress {getAddrVersion = v, getAddrData = d} = do
    hrp <- getBech32Prefix net
    segwitEncode hrp v (B.unpack d)

-- | Parse 'Base58', 'Bech32' or 'CashAddr' address, depending on network.
textToAddr :: Network -> Text -> Maybe Address
textToAddr net txt =
    cashToAddr net txt <|> bech32ToAddr net txt <|> base58ToAddr net txt

cashToAddr :: Network -> Text -> Maybe Address
cashToAddr net txt = do
    (ver, bs) <- cashAddrDecode net txt
    case ver of
        0 -> PubKeyAddress <$> eitherToMaybe (runGetS deserialize bs)
        1 -> ScriptAddress <$> eitherToMaybe (runGetS deserialize bs)
        _ -> Nothing

bech32ToAddr :: Network -> Text -> Maybe Address
bech32ToAddr net txt = do
    hrp <- getBech32Prefix net
    (ver, bs) <- second B.pack <$> segwitDecode hrp txt
    case ver of
        0 -> case B.length bs of
            20 -> WitnessPubKeyAddress <$> eitherToMaybe (runGetS deserialize bs)
            32 -> WitnessScriptAddress <$> eitherToMaybe (runGetS deserialize bs)
            _  -> Nothing
        _ -> Just $ WitnessAddress ver bs

base58ToAddr :: Network -> Text -> Maybe Address
base58ToAddr net txt =
    eitherToMaybe . runGetS (base58get net) =<< decodeBase58Check txt

base58get :: MonadGet m => Network -> m Address
base58get net = do
    pfx <- getWord8
    addr <- deserialize
    f pfx addr
  where
    f x a
        | x == getAddrPrefix net = return $ PubKeyAddress a
        | x == getScriptPrefix net = return $ ScriptAddress a
        | otherwise = fail "Does not recognize address prefix"

base58put :: MonadPut m => Network -> Address -> m ()
base58put net (PubKeyAddress h) = do
        putWord8 (getAddrPrefix net)
        serialize h
base58put net (ScriptAddress h) = do
        putWord8 (getScriptPrefix net)
        serialize h
base58put _ _ = error "Cannot serialize this address as Base58"

-- | Obtain a standard pay-to-public-key-hash address from a public key.
pubKeyAddr :: PubKeyI -> Address
pubKeyAddr = PubKeyAddress . addressHash . runPutS . serialize

-- | Obtain a standard pay-to-public-key-hash (P2PKH) address from a 'Hash160'.
p2pkhAddr :: Hash160 -> Address
p2pkhAddr = PubKeyAddress

-- | Obtain a SegWit pay-to-witness-public-key-hash (P2WPKH) address from a
-- public key.
pubKeyWitnessAddr :: PubKeyI -> Address
pubKeyWitnessAddr = WitnessPubKeyAddress . addressHash . runPutS . serialize

-- | Obtain a backwards-compatible SegWit P2SH-P2WPKH address from a public key.
pubKeyCompatWitnessAddr :: PubKeyI -> Address
pubKeyCompatWitnessAddr =
    p2shAddr .
    addressHash .
    encodeOutputBS .
    PayWitnessPKHash .
    addressHash .
    runPutS .
    serialize

-- | Obtain a SegWit pay-to-witness-public-key-hash (P2WPKH) address from a
-- 'Hash160'.
p2wpkhAddr :: Hash160 -> Address
p2wpkhAddr = WitnessPubKeyAddress

-- | Obtain a standard pay-to-script-hash (P2SH) address from a 'Hash160'.
p2shAddr :: Hash160 -> Address
p2shAddr = ScriptAddress

-- | Obtain a SegWit pay-to-witness-script-hash (P2WSH) address from a 'Hash256'
p2wshAddr :: Hash256 -> Address
p2wshAddr = WitnessScriptAddress

-- | Compute a standard pay-to-script-hash (P2SH) address for an output script.
payToScriptAddress :: ScriptOutput -> Address
payToScriptAddress = p2shAddr . addressHash . encodeOutputBS

-- | Compute a SegWit pay-to-witness-script-hash (P2WSH) address for an output
-- script.
payToWitnessScriptAddress :: ScriptOutput -> Address
payToWitnessScriptAddress = p2wshAddr . sha256 . encodeOutputBS

-- | Compute a backwards-compatible SegWit P2SH-P2WSH address.
payToNestedScriptAddress :: ScriptOutput -> Address
payToNestedScriptAddress =
    p2shAddr . addressHash . encodeOutputBS . toP2WSH . encodeOutput

-- | Encode an output script from an address. Will fail if using a
-- pay-to-witness address on a non-SegWit network.
addressToOutput :: Address -> ScriptOutput
addressToOutput =
    \case
        PubKeyAddress h        -> PayPKHash h
        ScriptAddress h        -> PayScriptHash h
        WitnessPubKeyAddress h -> PayWitnessPKHash h
        WitnessScriptAddress h -> PayWitnessScriptHash h
        WitnessAddress v d     -> PayWitness v d

-- | Get output script AST for an 'Address'.
addressToScript :: Address -> Script
addressToScript = encodeOutput . addressToOutput

-- | Encode address as output script in 'ByteString' form.
addressToScriptBS :: Address -> ByteString
addressToScriptBS = runPutS . serialize . addressToScript

-- | Decode an output script into an 'Address' if it has such representation.
scriptToAddress :: Script -> Either String Address
scriptToAddress =
    maybeToEither "Could not decode address" . outputAddress <=< decodeOutput

-- | Decode a serialized script into an 'Address'.
scriptToAddressBS :: ByteString -> Either String Address
scriptToAddressBS =
    maybeToEither "Could not decode address" . outputAddress <=< decodeOutputBS

-- | Get the 'Address' of a 'ScriptOutput'.
outputAddress :: ScriptOutput -> Maybe Address
outputAddress =
    \case
        PayPKHash h            -> Just $ PubKeyAddress h
        PayScriptHash h        -> Just $ ScriptAddress h
        PayPK k                -> Just $ pubKeyAddr k
        PayWitnessPKHash h     -> Just $ WitnessPubKeyAddress h
        PayWitnessScriptHash h -> Just $ WitnessScriptAddress h
        PayWitness v d         -> Just $ WitnessAddress v d
        _                      -> Nothing

-- | Infer the 'Address' of a 'ScriptInput'.
inputAddress :: ScriptInput -> Maybe Address
inputAddress =
    \case
        (RegularInput (SpendPKHash _ key)) -> Just $ pubKeyAddr key
        (ScriptHashInput _ rdm)            -> Just $ payToScriptAddress rdm
        _                                  -> Nothing