{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}

module Zenith where

import Control.Concurrent (threadDelay)
import Control.Monad
import Data.Aeson
import Data.Aeson.Types
import qualified Data.Array as A
import qualified Data.ByteString as B
import qualified Data.ByteString.Base64 as B64
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy as LB
import Data.Char
import Data.Functor (void)
import Data.HexString
import Data.Maybe
import qualified Data.Scientific as Scientific
import qualified Data.Text as T
import qualified Data.Text.Encoding as E
import Data.Text.Encoding.Error (lenientDecode)
import qualified Data.Text.IO as TIO
import qualified Data.Vector as V
import Data.Word
import GHC.Generics
import Haskoin.Address.Bech32
import Network.HTTP.Simple
import Network.HTTP.Types
import Numeric
import System.Clipboard
import System.Exit
import System.IO
import System.Process (createProcess_, shell)
import Text.Read (readMaybe)
import Text.Regex
import Text.Regex.Base
import Text.Regex.Posix

-- | A type to model Zcash RPC calls
data RpcCall =
  RpcCall
    { jsonrpc :: T.Text
    , id :: T.Text
    , method :: T.Text
    , params :: [Value]
    }
  deriving (Show, Generic, ToJSON, FromJSON)

-- | Type for modelling the different address sources for Zcash 5.0.0
data AddressSource
  = LegacyRandom
  | Imported
  | ImportedWatchOnly
  | KeyPool
  | LegacySeed
  | MnemonicSeed
  deriving (Read, Show, Eq, Generic, ToJSON)

instance FromJSON AddressSource where
  parseJSON =
    withText "AddressSource" $ \case
      "legacy_random" -> return LegacyRandom
      "imported" -> return Imported
      "imported_watchonly" -> return ImportedWatchOnly
      "keypool" -> return KeyPool
      "legacy_hdseed" -> return LegacySeed
      "mnemonic_seed" -> return MnemonicSeed
      _ -> fail "Not a known address source"

data ZcashPool
  = Transparent
  | Sprout
  | Sapling
  | Orchard
  deriving (Show, Eq, Generic, ToJSON)

instance FromJSON ZcashPool where
  parseJSON =
    withText "ZcashPool" $ \case
      "p2pkh" -> return Transparent
      "sprout" -> return Sprout
      "sapling" -> return Sapling
      "orchard" -> return Orchard
      _ -> fail "Not a known Zcash pool"

data ZcashAddress =
  ZcashAddress
    { source :: AddressSource
    , pool :: [ZcashPool]
    , account :: Maybe Integer
    , addy :: T.Text
    }
  deriving (Eq)

instance Show ZcashAddress where
  show (ZcashAddress s p i a) =
    T.unpack (T.take 8 a) ++
    "..." ++ T.unpack (T.takeEnd 8 a) ++ " Pools: " ++ show p

-- | A type to model the response of the Zcash RPC
data RpcResponse r =
  RpcResponse
    { err :: Maybe T.Text
    , respId :: T.Text
    , result :: r
    }
  deriving (Show, Generic, ToJSON)

instance (FromJSON r) => FromJSON (RpcResponse r) where
  parseJSON (Object obj) = do
    e <- obj .: "error"
    rId <- obj .: "id"
    r <- obj .: "result"
    pure $ RpcResponse e rId r
  parseJSON invalid =
    prependFailure
      "parsing RpcResponse failed, "
      (typeMismatch "Object" invalid)

newtype NodeVersion =
  NodeVersion Integer
  deriving (Eq, Show)

instance FromJSON NodeVersion where
  parseJSON =
    withObject "NodeVersion" $ \obj -> do
      v <- obj .: "version"
      pure $ NodeVersion v

-- | A type to model an address group
data AddressGroup =
  AddressGroup
    { agsource :: AddressSource
    , agtransparent :: [ZcashAddress]
    , agsapling :: [ZcashAddress]
    , agunified :: [ZcashAddress]
    }
  deriving (Show, Generic)

instance FromJSON AddressGroup where
  parseJSON =
    withObject "AddressGroup" $ \obj -> do
      s <- obj .: "source"
      t <- obj .:? "transparent"
      sap <- obj .:? "sapling"
      uni <- obj .:? "unified"
      sL <- processSapling sap s
      tL <- processTransparent t s
      uL <- processUnified uni
      return $ AddressGroup s tL (concat sL) (concat uL)
    where
      processTransparent c s1 =
        case c of
          Nothing -> return []
          Just x -> do
            x' <- x .: "addresses"
            return $ map (ZcashAddress s1 [Transparent] Nothing) x'
      processSapling k s2 =
        case k of
          Nothing -> return []
          Just y -> mapM (processOneSapling s2) y
            where processOneSapling sx =
                    withObject "Sapling" $ \oS -> do
                      oS' <- oS .: "addresses"
                      return $ map (ZcashAddress sx [Sapling] Nothing) oS'
      processUnified u =
        case u of
          Nothing -> return []
          Just z -> mapM processOneAccount z
            where processOneAccount =
                    withObject "UAs" $ \uS -> do
                      acct <- uS .: "account"
                      uS' <- uS .: "addresses"
                      mapM (processUAs acct) uS'
                    where
                      processUAs a =
                        withObject "UAs" $ \v -> do
                          addr <- v .: "address"
                          p <- v .: "receiver_types"
                          return $ ZcashAddress MnemonicSeed p a addr

displayZec :: Integer -> String
displayZec s
  | s < 100 = show s ++ " zats "
  | s < 100000 = show (fromIntegral s / 100) ++ " μZEC "
  | s < 100000000 = show (fromIntegral s / 100000) ++ " mZEC "
  | otherwise = show (fromIntegral s / 100000000) ++ " ZEC "

-- | A type to model a Zcash transaction
data ZcashTx =
  ZcashTx
    { ztxid :: T.Text
    , zamount :: Double
    , zamountZat :: Integer
    , zblockheight :: Integer
    , zblocktime :: Integer
    , zchange :: Bool
    , zconfirmations :: Integer
    , zmemo :: T.Text
    }
  deriving (Show, Generic)

instance FromJSON ZcashTx where
  parseJSON =
    withObject "ZcashTx" $ \obj -> do
      t <- obj .: "txid"
      a <- obj .: "amount"
      aZ <- obj .: "amountZat"
      bh <- obj .: "blockheight"
      bt <- obj .: "blocktime"
      c <- obj .:? "change"
      conf <- obj .: "confirmations"
      m <- obj .:? "memo"
      pure $
        ZcashTx
          t
          a
          aZ
          bh
          bt
          (fromMaybe False c)
          conf
          (case m of
             Nothing -> ""
             Just m' -> T.filter (/= '\NUL') $ decodeHexText m')

instance ToJSON ZcashTx where
  toJSON (ZcashTx t a aZ bh bt c conf m) =
    object
      [ "amount" .= a
      , "amountZat" .= aZ
      , "txid" .= t
      , "blockheight" .= bh
      , "blocktime" .= bt
      , "change" .= c
      , "confirmations" .= conf
      , "memo" .= m
      ]

-- | Type for the UA balance
data UABalance =
  UABalance
    { uatransparent :: Integer
    , uasapling :: Integer
    , uaorchard :: Integer
    }
  deriving (Eq)

instance Show UABalance where
  show (UABalance t s o) =
    " T: " ++ show t ++ " S: " ++ show s ++ " O: " ++ show o

instance FromJSON UABalance where
  parseJSON =
    withObject "UABalance" $ \obj -> do
      p <- obj .: "pools"
      t <- p .:? "transparent"
      s <- p .:? "sapling"
      o <- p .:? "orchard"
      vT <-
        case t of
          Nothing -> return 0
          Just t' -> t' .: "valueZat"
      vS <-
        case s of
          Nothing -> return 0
          Just s' -> s' .: "valueZat"
      vO <-
        case o of
          Nothing -> return 0
          Just o' -> o' .: "valueZat"
      pure $ UABalance vT vS vO

-- | Type for Operation Result
data OpResult =
  OpResult
    { opsuccess :: T.Text
    , opmessage :: Maybe T.Text
    , optxid :: Maybe T.Text
    }
  deriving (Show, Eq)

instance FromJSON OpResult where
  parseJSON =
    withObject "OpResult" $ \obj -> do
      s <- obj .: "status"
      r <- obj .:? "result"
      e <- obj .:? "error"
      t <-
        case r of
          Nothing -> return Nothing
          Just r' -> r' .: "txid"
      m <-
        case e of
          Nothing -> return Nothing
          Just m' -> m' .: "message"
      pure $ OpResult s m t

-- | Helper function to turn a hex-encoded memo strings to readable text
decodeHexText :: String -> T.Text
decodeHexText h = E.decodeUtf8With lenientDecode $ B.pack $ hexRead h
  where
    hexRead hexText
      | null chunk = []
      | otherwise =
        fromIntegral (read ("0x" <> chunk)) : hexRead (drop 2 hexText)
      where
        chunk = take 2 hexText

-- | Helper function to turn a string into a hex-encoded string
encodeHexText :: String -> String
encodeHexText t = mconcat (map padHex t)
  where
    padHex x =
      if ord x < 16
        then "0" ++ (showHex . ord) x ""
        else showHex (ord x) ""

encodeHexText' :: T.Text -> String
encodeHexText' t =
  if T.length t > 0
    then T.unpack . toText . fromBytes $ E.encodeUtf8 t
    else T.unpack . toText . fromBytes $ E.encodeUtf8 "Sent from Zenith"

-- | Helper function to extract addresses from AddressGroups
getAddresses :: AddressGroup -> [ZcashAddress]
getAddresses ag = agtransparent ag <> agsapling ag <> agunified ag

-- | Helper function to validate potential Zcash addresses
validateAddress :: T.Text -> Maybe ZcashPool
validateAddress txt --(tReg || sReg && isJust chk) || (uReg && isJust chk)
  | tReg = Just Transparent
  | sReg && isJust chk = Just Sapling
  | uReg && isJust chk = Just Orchard
  | otherwise = Nothing
  where
    transparentRegex = "^t1[a-zA-Z0-9]{33}$" :: String
    shieldedRegex = "^zs[a-zA-Z0-9]{76}$" :: String
    unifiedRegex = "^u[a-zA-Z0-9]" :: String
    tReg = T.unpack txt =~ transparentRegex :: Bool
    sReg = T.unpack txt =~ shieldedRegex :: Bool
    uReg = T.unpack txt =~ unifiedRegex :: Bool
    chk = bech32mDecode txt

-- | RPC methods
-- | List addresses
listAddresses :: B.ByteString -> B.ByteString -> IO [ZcashAddress]
listAddresses user pwd = do
  response <- makeZcashCall user pwd "listaddresses" []
  let rpcResp = decode response :: Maybe (RpcResponse [AddressGroup])
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just res -> do
      let addys = result res
      let addList = concatMap getAddresses addys
      return addList

-- | Get address balance
getBalance :: B.ByteString -> B.ByteString -> ZcashAddress -> IO [Integer]
getBalance user pwd zadd = do
  let a = account zadd
  case a of
    Nothing -> do
      response <-
        makeZcashCall
          user
          pwd
          "z_getbalance"
          [ String (addy zadd)
          , Number (Scientific.scientific 1 0)
          , Data.Aeson.Bool True
          ]
      let rpcResp = decode response :: Maybe (RpcResponse Integer)
      case rpcResp of
        Nothing -> fail "Couldn't parse node response"
        Just res -> do
          return [result res]
    Just acct -> do
      response <-
        makeZcashCall
          user
          pwd
          "z_getbalanceforaccount"
          [Number (Scientific.scientific acct 0)]
      let rpcResp = decode response :: Maybe (RpcResponse UABalance)
      case rpcResp of
        Nothing -> fail "Couldn't parse node response"
        Just res -> do
          return $ readUABalance (result res)
          where readUABalance ua =
                  [uatransparent ua, uasapling ua, uaorchard ua]

-- | List transactions
listTxs :: B.ByteString -> B.ByteString -> ZcashAddress -> IO [ZcashTx]
listTxs user pwd zaddy = do
  response <-
    makeZcashCall user pwd "z_listreceivedbyaddress" [String $ addy zaddy]
  let rpcResp = decode response :: Maybe (RpcResponse [ZcashTx])
  case rpcResp of
    Nothing -> fail "listTxs: Couldn't parse node response"
    Just res -> do
      return $ result res

-- | Send Tx
sendTx ::
     B.ByteString
  -> B.ByteString
  -> ZcashAddress
  -> T.Text
  -> Double
  -> Maybe T.Text
  -> IO ()
sendTx user pwd fromAddy toAddy amount memo = do
  bal <- getBalance user pwd fromAddy
  let valAdd = validateAddress toAddy
  if sum bal - floor (amount * 100000000) >= 1000
    then do
      if source fromAddy /= ImportedWatchOnly
        then do
          let privacyPolicy
                | valAdd == Just Transparent = "AllowRevealedRecipients"
                | isNothing (account fromAddy) &&
                    elem Transparent (pool fromAddy) = "AllowRevealedSenders"
                | otherwise = "AllowRevealedAmounts"
          let pd =
                case memo of
                  Nothing ->
                    [ Data.Aeson.String (addy fromAddy)
                    , Data.Aeson.Array
                        (V.fromList
                           [object ["address" .= toAddy, "amount" .= amount]])
                    , Data.Aeson.Number $ Scientific.scientific 1 1
                    , Data.Aeson.Number $ Scientific.scientific 1 (-5)
                    , Data.Aeson.String privacyPolicy
                    ]
                  Just memo' ->
                    [ Data.Aeson.String (addy fromAddy)
                    , Data.Aeson.Array
                        (V.fromList
                           [ object
                               [ "address" .= toAddy
                               , "amount" .= amount
                               , "memo" .= encodeHexText' memo'
                               ]
                           ])
                    , Data.Aeson.Number $ Scientific.scientific 1 1
                    , Data.Aeson.Number $ Scientific.scientific 1 (-5)
                    , Data.Aeson.String privacyPolicy
                    ]
          response <- makeZcashCall user pwd "z_sendmany" pd
          let rpcResp = decode response :: Maybe (RpcResponse T.Text)
          case rpcResp of
            Nothing -> fail "Couldn't parse node response"
            Just res -> do
              putStr "  Sending."
              checkOpResult user pwd (result res)
        else putStrLn "Error: Source address is view-only."
    else putStrLn "Error: Insufficient balance in source address."

-- | Make a Zcash RPC call
makeZcashCall ::
     B.ByteString
  -> B.ByteString
  -> T.Text
  -> [Data.Aeson.Value]
  -> IO LB.ByteString
makeZcashCall username password m p = do
  let payload = RpcCall "1.0" "test" m p
  let myRequest =
        setRequestBodyJSON payload $
        setRequestPort 8232 $
        setRequestBasicAuth username password $
        setRequestMethod "POST" defaultRequest
  response <- httpLBS myRequest
  let respStatus = getResponseStatusCode response
  let body = getResponseBody response
  case respStatus of
    500 -> do
      let rpcResp = decode body :: Maybe (RpcResponse String)
      case rpcResp of
        Nothing -> fail $ "Unknown server error " ++ show response
        Just x -> fail (result x)
    401 -> fail "Incorrect full node credentials"
    200 -> return body
    _ -> fail "Unknown error"

-- | Display an address
displayZcashAddress ::
     B.ByteString -> B.ByteString -> (Int, ZcashAddress) -> IO ()
displayZcashAddress user pwd (idx, zaddy) = do
  zats <- getBalance user pwd zaddy
  putStr $ show idx ++ ": "
  putStr $ show zaddy
  when (source zaddy == ImportedWatchOnly) (putStr "[VK]")
  putStr " Balance: "
  mapM_ (putStr . displayZec) zats
  putStrLn ""

-- | Copy an address to the clipboard
copyAddress :: ZcashAddress -> IO ()
copyAddress a =
  void $
  createProcess_ "toClipboard" $
  shell $ "echo " ++ T.unpack (addy a) ++ " | xclip -r -selection clipboard"

-- | Verify operation result
checkOpResult :: B.ByteString -> B.ByteString -> T.Text -> IO ()
checkOpResult user pwd opid = do
  response <-
    makeZcashCall
      user
      pwd
      "z_getoperationstatus"
      [Data.Aeson.Array (V.fromList [Data.Aeson.String opid])]
  let rpcResp = decode response :: Maybe (RpcResponse [OpResult])
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just res -> do
      let r = result res
      mapM_ showResult r
  where
    showResult t =
      case opsuccess t of
        "success" ->
          putStrLn $ "  Success! Tx ID: " ++ maybe "" T.unpack (optxid t)
        "executing" -> do
          putStr "."
          hFlush stdout
          threadDelay 1000000 >> checkOpResult user pwd opid
        _ -> putStrLn $ "  Failed :( " ++ maybe "" T.unpack (opmessage t)

-- | Check for accounts
checkAccounts :: B.ByteString -> B.ByteString -> IO Bool
checkAccounts user pwd = do
  response <- makeZcashCall user pwd "z_listaccounts" []
  let rpcResp = decode response :: Maybe (RpcResponse [Object])
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just res -> do
      let r = result res
      return $ not (null r)

-- | Add account to node
createAccount :: B.ByteString -> B.ByteString -> IO ()
createAccount user pwd = do
  response <- makeZcashCall user pwd "z_getnewaccount" []
  let rpcResp = decode response :: Maybe (RpcResponse Object)
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just res -> do
      let r = result res
      putStrLn "  Account created!"

-- | Create new Unified Address
createUnifiedAddress :: B.ByteString -> B.ByteString -> Bool -> Bool -> IO ()
createUnifiedAddress user pwd tRec sRec = do
  let recs = getReceivers tRec sRec
  let pd = [Data.Aeson.Number $ Scientific.scientific 0 1, recs]
  newResp <- makeZcashCall user pwd "z_getaddressforaccount" pd
  let rpcResp = decode newResp :: Maybe (RpcResponse Object)
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just res -> do
      let r = result res
      putStrLn "  New UA created!"
  where
    getReceivers t s
      | t && s =
        Data.Aeson.Array
          (V.fromList
             [ Data.Aeson.String "p2pkh"
             , Data.Aeson.String "sapling"
             , Data.Aeson.String "orchard"
             ])
      | t =
        Data.Aeson.Array
          (V.fromList [Data.Aeson.String "p2pkh", Data.Aeson.String "orchard"])
      | s =
        Data.Aeson.Array
          (V.fromList [Data.Aeson.String "sapling", Data.Aeson.String "orchard"])
      | otherwise = Data.Aeson.Array (V.fromList [Data.Aeson.String "orchard"])

-- | Check Zcash full node server
checkServer :: B.ByteString -> B.ByteString -> IO ()
checkServer user pwd = do
  resp <- makeZcashCall user pwd "getinfo" []
  let rpcResp = decode resp :: Maybe (RpcResponse NodeVersion)
  case rpcResp of
    Nothing -> fail "Couldn't parse node response"
    Just myResp -> do
      let r = result myResp
      if isNodeValid r
        then putStrLn $ "Connected to Zcash Full Node (" <> show r <> ") :)"
        else do
          putStrLn "Deprecated Zcash Full Node version found. Exiting"
          exitFailure
      where isNodeValid (NodeVersion i) = i >= 5000000

-- | Read ZIP-321 URI
sendWithUri ::
     B.ByteString -> B.ByteString -> ZcashAddress -> String -> Bool -> IO ()
sendWithUri user pwd fromAddy uri repTo = do
  let uriRegex = mkRegex "^zcash:(\\w+)\\?amount=(.*)\\&memo=(.*)$"
  if matchTest uriRegex uri
    then do
      let reg = matchAllText uriRegex uri
      let parsedAddress = fst $ head reg A.! 1
      let parsedAmount = fst $ head reg A.! 2
      let parsedEncodedMemo = fst $ head reg A.! 3
      let addType = validateAddress $ T.pack parsedAddress
      case addType of
        Nothing -> putStrLn "   Invalid address"
        Just Transparent -> do
          putStrLn $ "   Address is valid: " ++ parsedAddress
          case (readMaybe parsedAmount :: Maybe Double) of
            Nothing -> putStrLn "   Invalid amount."
            Just amt -> do
              putStrLn $ "   Valid ZEC amount: " ++ show amt
              sendTx user pwd fromAddy (T.pack parsedAddress) amt Nothing
        Just _ -> do
          putStrLn $ "   Address is valid: " ++ parsedAddress
          case (readMaybe parsedAmount :: Maybe Double) of
            Nothing -> putStrLn "   Invalid amount."
            Just amt -> do
              putStrLn $ "   Valid ZEC amount: " ++ show amt
              let decodedMemo =
                    E.decodeUtf8With lenientDecode $
                    B64.decodeLenient $ C.pack parsedEncodedMemo
              TIO.putStrLn $ "   Memo: " <> decodedMemo
              sendTx
                user
                pwd
                fromAddy
                (T.pack parsedAddress)
                amt
                (if repTo
                   then Just $
                        T.concat [decodedMemo, "\nReply-To:\n", addy fromAddy]
                   else Just decodedMemo)
    else putStrLn "URI is not compliant with ZIP-321"