-- Copyright 2022-2024 Vergara Technologies LLC
-- 
-- This file is part of Zcash-Haskell.
--
-- |
-- Module      :  ZcashHaskell.Sapling
-- Copyright   :  2022-2024 Vergara Technologies
-- License     :  MIT
--
-- Maintainer  :  pitmutt@vergara.tech
-- Stability   :  experimental
-- Portability :  unknown
--
-- Functions to interact with the Sapling shielded pool of the Zcash blockchain.
--
{-# LANGUAGE OverloadedStrings #-}

module ZcashHaskell.Sapling where

import C.Zcash
  ( rustWrapperIsShielded
  , rustWrapperReadSaplingCommitmentTree
  , rustWrapperReadSaplingPosition
  , rustWrapperReadSaplingWitness
  , rustWrapperDecodeSaplingAddress
  , rustWrapperSaplingCheck
  , rustWrapperSaplingChgPaymentAddress
  , rustWrapperSaplingDecodeEsk
  , rustWrapperSaplingNoteDecode
  , rustWrapperSaplingPaymentAddress
  , rustWrapperSaplingSpendingkey
  , rustWrapperSaplingVkDecode
  , rustWrapperTxParse
  )
import Data.Aeson
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C
import qualified Data.Text as T
import Data.HexString (HexString(..), fromText, hexString, toBytes, toText)
import Data.Word
import Foreign.Rust.Marshall.Variable
  ( withPureBorshVarBuffer
  , withPureBorshVarBuffer
  )
import ZcashHaskell.Types  
import ZcashHaskell.Utils (decodeBech32, encodeBech32, encodeBech32m)

-- | Check if given bytesting is a valid encoded shielded address
isValidShieldedAddress :: BS.ByteString -> Bool
isValidShieldedAddress = rustWrapperIsShielded

getShieldedOutputs :: HexString -> [BS.ByteString]
getShieldedOutputs t = withPureBorshVarBuffer $ rustWrapperTxParse $ toBytes t

serializeShieldedOutput :: ShieldedOutput -> BS.ByteString
serializeShieldedOutput so =
  hexBytes . fromText $
  toText (s_cv so) <>
  toText (s_cmu so) <>
  toText (s_ephKey so) <>
  toText (s_encCipherText so) <>
  toText (s_outCipherText so) <> toText (s_proof so)

-- | Check if given bytestring is a valid Sapling viewing key
isValidSaplingViewingKey :: BS.ByteString -> Bool
isValidSaplingViewingKey k =
  case hrp decodedKey of
    "zxviews" -> rustWrapperSaplingVkDecode $ bytes decodedKey
    _ -> False
  where
    decodedKey = decodeBech32 k

-- | Check if the given bytestring for the Sapling viewing key matches the second bytestring for the address
matchSaplingAddress :: BS.ByteString -> BS.ByteString -> Bool
matchSaplingAddress = rustWrapperSaplingCheck

-- | Attempt to decode the given raw tx with the given Sapling viewing key
decodeSaplingOutput :: BS.ByteString -> BS.ByteString -> Maybe DecodedNote
decodeSaplingOutput key out =
  case a_value decodedAction of
    0 -> Nothing
    _ -> Just decodedAction
  where
    decodedAction =
      withPureBorshVarBuffer $ rustWrapperSaplingNoteDecode key out

instance FromJSON RawTxResponse where
  parseJSON =
    withObject "RawTxResponse" $ \obj -> do
      i <- obj .: "txid"
      o <- obj .:? "orchard"
      h <- obj .: "hex"
      ht <- obj .: "height"
      c <- obj .: "confirmations"
      b <- obj .: "blocktime"
      sSpend <- obj .: "vShieldedSpend"
      case o of
        Nothing ->
          pure $ RawTxResponse i h sSpend (getShieldedOutputs h) [] ht c b
        Just o' -> do
          a <- o' .: "actions"
          pure $ RawTxResponse i h sSpend (getShieldedOutputs h) a ht c b

-- | Attempt to decode the given raw tx with the given Sapling spending key
decodeSaplingOutputEsk ::
     SaplingSpendingKey
  -> ShieldedOutput
  -> ZcashNet
  -> Scope
  -> Integer
  -> Maybe DecodedNote
decodeSaplingOutputEsk key out znet scope pos =
  case a_value decodedAction of
    0 -> Nothing
    _ -> Just decodedAction
  where
    decodedAction =
      withPureBorshVarBuffer $
      rustWrapperSaplingDecodeEsk
        (getBytes key)
        (serializeShieldedOutput out)
        (scope == External)
        (znet == MainNet)
        (fromIntegral pos)

-- | Attempts to obtain a sapling SpendingKey using a HDSeed
genSaplingSpendingKey :: Seed -> CoinType -> Int -> Maybe SaplingSpendingKey
genSaplingSpendingKey seed c i = do
  if BS.length res == 169
    then Just $ SaplingSpendingKey res
    else Nothing
  where
    res =
      withPureBorshVarBuffer
        (rustWrapperSaplingSpendingkey
           (getBytes seed)
           (fromIntegral $ getValue c)
           (fromIntegral i))

-- | Attempts to generate a sapling Payment Address using an ExtendedSpendingKey and a Diversifier Index 
genSaplingPaymentAddress :: Int -> SaplingSpendingKey -> Maybe SaplingReceiver
genSaplingPaymentAddress i extspk =
  if BS.length res == 43
    then Just $ SaplingReceiver res
    else Nothing
  where
    res =
      withPureBorshVarBuffer
        (rustWrapperSaplingPaymentAddress
           (getBytes extspk)
           (fromIntegral (i * 111)))

-- | Generate an internal Sapling address
genSaplingInternalAddress :: SaplingSpendingKey -> Maybe SaplingReceiver
genSaplingInternalAddress sk =
  if BS.length res == 43
    then Just $ SaplingReceiver res
    else Nothing
  where
    res =
      withPureBorshVarBuffer (rustWrapperSaplingChgPaymentAddress $ getBytes sk)

-- | Update a Sapling commitment tree
updateSaplingCommitmentTree ::
     SaplingCommitmentTree -- ^ the base tree
  -> HexString -- ^ the new note commitment
  -> Maybe SaplingCommitmentTree
updateSaplingCommitmentTree tree cmu =
  if BS.length (hexBytes updatedTree) > 1
    then Just $ SaplingCommitmentTree updatedTree
    else Nothing
  where
    updatedTree =
      withPureBorshVarBuffer $
      rustWrapperReadSaplingCommitmentTree
        (hexBytes $ sapTree tree)
        (hexBytes cmu)

-- | Get the Sapling incremental witness from a commitment tree
getSaplingWitness :: SaplingCommitmentTree -> Maybe SaplingWitness
getSaplingWitness tree =
  if BS.length (hexBytes wit) > 1
    then Just $ SaplingWitness wit
    else Nothing
  where
    wit =
      withPureBorshVarBuffer $
      rustWrapperReadSaplingWitness (hexBytes $ sapTree tree)

-- | Get the Sapling note position from a witness
getSaplingNotePosition :: SaplingWitness -> Integer
getSaplingNotePosition =
  fromIntegral . rustWrapperReadSaplingPosition . hexBytes . sapWit

-- | Encode a SaplingReceiver into HRF text
encodeSaplingAddress :: ZcashNet -> SaplingReceiver -> Maybe T.Text
encodeSaplingAddress net sr = do 
  case net of
    MainNet ->
      Just $ encodeBech32 (C.pack sapPaymentAddressHrp) (getBytes sr)
    TestNet ->
      Just $ encodeBech32 (C.pack sapTestPaymentAddressHrp) (getBytes sr)

-- | Helper to get de Nework Id from FFI response
getNetId:: [Word8] -> ZcashNet
getNetId [x] = do 
  case x of 
    1 -> MainNet
    2 -> TestNet 

-- | decode a Sapling address 
decodeSaplingAddress :: BS.ByteString -> Maybe SaplingAddress
decodeSaplingAddress sapling_address = do 
  if BS.length sa > 1
    then do
      let sa0 = BS.unpack sa
      Just $ SaplingAddress  (getNetId (take 1 sa0)) 
           $ SaplingReceiver (BS.pack (drop 1 sa0))
    else Nothing
  where 
    sa = 
      withPureBorshVarBuffer $
      rustWrapperDecodeSaplingAddress sapling_address