{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoFieldSelectors #-}

module Main (main) where

import Control.Monad (foldM, guard)
import Data.Bifunctor
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BSB
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy as BSL
import Data.Fixed
import qualified Data.Map.Strict as Map
import Data.Maybe (isJust, isNothing)
import Data.Monoid
import Data.Word (Word64)
import GHC.Generics
import qualified Sel
import qualified Sel.Hashing.SHA512 as Hash
import qualified Sel.PublicKey.Signature as Sign
import Test.Hspec

main :: IO ()
main = Sel.secureMain $ hspec spec

spec :: Spec
spec = do
  describe "Cryptography Helpers" $ do
    it "verifies a correctly signed transaction" $ do
      (pubKey, secKey) <- generateKeyPair
      let txId = TxId 1
      sig <- signTx secKey txId
      verifySignature pubKey sig txId `shouldBe` True

    it "rejects a signature for a modified transaction ID" $ do
      (pubKey, secKey) <- generateKeyPair
      let txId1 = TxId 1
          txId2 = TxId 2
      sig <- signTx secKey txId1
      verifySignature pubKey sig txId2 `shouldBe` False

    it "rejects a valid signature verified against the wrong public key" $ do
      (_pubKey1, secKey1) <- generateKeyPair
      (pubKey2, _secKey2) <- generateKeyPair
      let txId = TxId 1
      sig <- signTx secKey1 txId
      verifySignature pubKey2 sig txId `shouldBe` False

    it "derives consistent addresses from public keys" $ do
      (pubKey, _) <- generateKeyPair
      deriveAddress pubKey `shouldBe` deriveAddress pubKey

  describe "Transaction Application (applyTx)" $ do
    it "successfully applies a valid transaction" $ do
      (pubKey, secKey) <- generateKeyPair
      let addr = deriveAddress pubKey
          txIdOrigin = TxId 0
          outRef = TxOutRef txIdOrigin (TxOutIdx 0)
          utxo = UTxO $ Map.singleton outRef (TxOut addr (Ada 100))

          txIdNew = TxId 1

      sig <- signTx secKey txIdNew
      let txIn = TxIn outRef pubKey sig -- TxIn now requires the PublicKey to verify against the Address
          txOut = TxOut addr (Ada 100)
          tx = Transaction txIdNew [txIn] [txOut]

      applyTx tx utxo `shouldSatisfy` isJust

    it "fails to apply a transaction with an invalid signature" $ do
      (pubKey1, _secKey1) <- generateKeyPair
      (_pubKey2, secKey2) <- generateKeyPair
      let addr1 = deriveAddress pubKey1
          txIdOrigin = TxId 0
          outRef = TxOutRef txIdOrigin (TxOutIdx 0)
          utxo = UTxO $ Map.singleton outRef (TxOut addr1 (Ada 100))

          txIdNew = TxId 1

      -- Signed by key 2, but claiming pubKey1 to spend addr1's funds
      sig <- signTx secKey2 txIdNew
      let txIn = TxIn outRef pubKey1 sig
          txOut = TxOut addr1 (Ada 100)
          tx = Transaction txIdNew [txIn] [txOut]

      applyTx tx utxo `shouldSatisfy` isNothing

-- The UTxO set maps these references to the actual locked outputs.
newtype UTxO = UTxO
  {unUTxO :: Map.Map TxOutRef TxOut}
  deriving newtype (Eq, Show)

-- TxOut simply locks a value (Ada) to an Address. It doesn't know about wallets.
data TxOut = TxOut
  { address :: Address Wallet,
    value :: Ada
  }
  deriving (Eq, Show)

-- An output reference uniquely identifies a specific unspent output.
data TxOutRef = TxOutRef
  { txId :: TxId,
    txOutIdx :: TxOutIdx
  }
  deriving (Eq, Ord, Show)

-- Will be the hash of the transaction
newtype TxId
  = TxId {unTxId :: Word64}
  deriving newtype (Eq, Ord, Show)

-- The index of the output of a transaction
newtype TxOutIdx
  = TxOutIdx {unTxOutIdx :: Int}
  deriving newtype (Eq, Ord, Show, Num, Enum)

newtype Ada
  = Ada {unAda :: Fixed E6}
  deriving newtype (Eq, Ord, Show)
  deriving (Semigroup, Monoid) via (Sum (Fixed E6))

-- TxIn points to the TxOutRef and provides the signature (Redeem) to unlock it.
data TxIn = TxIn
  { txOutRef :: TxOutRef,
    pubKey :: PublicKey,
    redeem :: Signature
  }
  deriving (Eq, Show)

-- A Transaction is simply a list of inputs being consumed and new outputs being created.
data Transaction = Transaction
  { txId :: TxId, -- Will be the hash of this transaction
    txInputs :: [TxIn],
    txOutputs :: [TxOut]
  }
  deriving stock (Eq, Show, Generic)

applyTx :: Transaction -> UTxO -> Maybe UTxO
applyTx tx (UTxO utxo) = do
  -- 1. Lookup all consumed outputs and pair them with their inputs
  consumed <- mapM (\txIn -> (txIn,) <$> utxo Map.!? txIn.txOutRef) tx.txInputs

  -- 2. Verify authorization
  guard $ all (\(txIn, txOut) -> verifySignatureRedeem txOut.address txIn.pubKey txIn.redeem tx.txId) consumed

  -- 3. Verify conservation of value
  let inputsValue = foldMap ((.value) . snd) consumed
      outputsValue = foldMap (.value) tx.txOutputs
  guard $ inputsValue == outputsValue

  -- 4. Update the UTxO map
  let utxoMinusInputs = foldl' (flip Map.delete) utxo ((.txOutRef) <$> tx.txInputs)
  return $
    UTxO $
      foldl'
        (\u (idx, out) -> Map.insert (TxOutRef tx.txId idx) out u)
        utxoMinusInputs
        (zip [0 ..] tx.txOutputs)

-- | Verifies that the provided PublicKey matches the locking Address,
-- and that the Signature is valid for the TxId.
verifySignatureRedeem :: Address Wallet -> PublicKey -> Signature -> TxId -> Bool
verifySignatureRedeem addr pubKey sig txId =
  verifyAddress pubKey addr && verifySignature pubKey sig txId

data Wallet = Wallet
  { key :: PublicKey,
    balance :: Ada
  }
  deriving stock (Eq, Show, Generic)

newtype WalletId
  = WalletId {unWalletId :: Address Wallet}
  deriving newtype (Eq, Ord, Show)

newtype Address a
  = Address {unAddress :: Hash.Hash}
  deriving newtype (Eq, Ord, Show)

newtype PublicKey
  = PublicKey {unPublicKey :: Sign.PublicKey}
  deriving newtype (Eq)

instance Show PublicKey where
  show _ = "<PublicKey>"

data Block = Block
  { blId :: BlockId,
    blTxs :: [Transaction]
  }
  deriving stock (Eq, Show, Generic)

newtype BlockId
  = BlockId {unBlockId :: Int}
  deriving newtype (Eq, Ord, Show)

applyBlock :: Block -> UTxO -> Maybe UTxO
applyBlock block utxo = foldM (flip applyTx) utxo block.blTxs

-- A basic cryptographic signature
newtype Signature
  = Signature {unSignaure :: Sign.SignedMessage}
  deriving newtype (Eq)

instance Show Signature where
  show _ = "<Signature>"

-- | Hashes a public key to produce a standard Wallet Address.
deriveAddress :: PublicKey -> Address a
deriveAddress pub = Address $ Hash.hashByteString (Sign.publicKeyToBinary pub.unPublicKey)

-- | Signs a payload (the TxId) using the sender's Ed25519 SecretKey.
signTx :: SecretKey -> TxId -> IO Signature
signTx sek txid = Signature <$> Sign.signMessage (bsTxId txid) sek.unSecretKey

-- | Verifies an Ed25519 signature against a PublicKey and the expected TxId payload.
verifySignature :: PublicKey -> Signature -> TxId -> Bool
verifySignature pub sig txid = Sign.openMessage sig.unSignaure pub.unPublicKey == Just (bsTxId txid)

-- | Checks if a given PublicKey mathematically hashes to the provided Address.
verifyAddress :: PublicKey -> Address a -> Bool
verifyAddress pub addr = deriveAddress pub == addr

newtype SecretKey
  = SecretKey {unSecretKey :: Sign.SecretKey}
  deriving newtype (Eq)

instance Show SecretKey where
  show _ = "<SecretKey>"

-- | Generates a new Ed25519 key pair using Sel.
generateKeyPair :: IO (PublicKey, SecretKey)
generateKeyPair = bimap PublicKey SecretKey <$> Sign.generateKeyPair

bsTxId :: TxId -> BS.StrictByteString
bsTxId = BSL.toStrict . BSB.toLazyByteString . BSB.word64Hex . (.unTxId)