module Crypto.XEd25519 where

import Control.Arrow
import Data.Bits
import Data.ByteArray as BA
import Data.Memory.PtrMethods (memCopy)
import Crypto.Hash
import Crypto.ECC.Edwards25519
import Crypto.Error
import qualified Crypto.PubKey.Ed25519 as Ed25519
import Foreign.Marshal
import Foreign.Ptr
import Foreign.Storable
import qualified Crypto.PubKey.Curve25519 as X25519

import Crypto.XEd25519.FieldElement
-- import Crypto.Nonce


data SecretKey = SecretKey { secretScalar :: Scalar }

data PublicKey = PublicKey Ed25519.PublicKey
 deriving Eq

-- type Nonce = Nonce32
type Nonce = Bytes -- 32 bytes

newtype EncodedPoint = EncodedPoint Point

instance ByteArrayAccess SecretKey where
    length _ = 32
    withByteArray (SecretKey scalar) = withByteArray (scalarEncode scalar :: Bytes)

instance ByteArrayAccess PublicKey where
    length _ = 32
    withByteArray (PublicKey edpub) = withByteArray edpub

instance ByteArrayAccess EncodedPoint where
    length _ = 32
    withByteArray (EncodedPoint pt) f =
        withByteArray (pointEncode pt :: Bytes) f


data Signature = Signature EncodedPoint Scalar

instance ByteArrayAccess Signature where
    length _ = 64
    withByteArray (Signature pt scalar) f =
        withByteArray pt                 $ \ptptr -> do
        withByteArray (SecretKey scalar) $ \scalarptr -> do
        allocaBytes 64                   $ \ptr -> do
            memCopy ptr ptptr 32
            memCopy (ptr `plusPtr` 32) scalarptr 32
            f (castPtr ptr)

signatureDecode :: ByteArrayAccess ba => ba -> Maybe Signature
signatureDecode bs = do
    let (ptbs,scbs) = BA.splitAt 32 ( BA.convert bs :: Bytes)
    pt <- maybeCryptoError $ pointDecode ptbs
    sc <- maybeCryptoError $ scalarDecodeLong scbs
    return $ Signature (ge_p3_tobytes pt) sc


padding :: Bytes
padding = 0xFE `BA.cons` BA.replicate 31 0xFF

sign :: ByteArrayAccess dta => dta -> Nonce -> SecretKey -> PublicKey -> Signature
sign dta nonce sec pub = Signature rB s
 where
    rB = ge_p3_tobytes $ ge_scalarmult_base r

    r = sc_reduce $ hashFinalize $ (`hashUpdate` padding)
                                    >>> (`hashUpdate` sec)
                                    >>> (`hashUpdate` dta)
                                    >>> (`hashUpdate` nonce) $ hashInit

    h = sc_reduce $ hashFinalize $ (`hashUpdate` rB)
                                    >>> (`hashUpdate` pub)
                                    >>> (`hashUpdate` dta) $ hashInit

    -- s = r + ha (mod q)
    s = sc_muladd h (secretScalar sec) r



ge_p3_tobytes :: Point -> EncodedPoint
ge_p3_tobytes = EncodedPoint

ge_scalarmult_base :: Scalar -> Point
ge_scalarmult_base = toPoint

sc_muladd :: Scalar -> Scalar -> Scalar -> Scalar
sc_muladd a b c = scalarAdd (scalarMul a b) c

sc_reduce :: Digest SHA512 -> Scalar
sc_reduce digest = x where CryptoPassed x = scalarDecodeLong digest -- ???

-- Scalar is internally, at least on 64bit machines, represented as 5
-- 56-bit words in little-endian order, each encoded as a Word64.
sc_neg :: Scalar -> Scalar
sc_neg = scalarMul sc_neg1

verify :: ByteArrayAccess dta => PublicKey -> dta -> Signature -> Bool
verify pub dta signature = Ed25519.verify ed_pub dta ed_sig
 where
    CryptoPassed ed_pub = Ed25519.publicKey pub'
    CryptoPassed ed_sig = Ed25519.signature signature'

    -- Get the sign bit from the s part of the signature.
    sign_bit = BA.index signature 63 .&. 0x80

    -- Set the sign bit to zero in the s part of the signature.
    signature' :: Bytes
    signature' = BA.copyAndFreeze signature $ \ptr -> do
        let at63 = plusPtr ptr 63
        byte63 <- peek at63
        poke at63 $ byte63 .&. (0x7F `asTypeOf` sign_bit)

    -- Restore the sign bit on the verification key, which should have 0 as its
    -- current sign bit.
    pub' :: Bytes
    pub' = BA.copyAndFreeze pub $ \ptr -> do
        let at31 = plusPtr ptr 31
        byte31 <- peek at31
        poke at31 $ (byte31 .&. 0x7F) .|. sign_bit


-- typedef crypto_int32 fe[10];
--
-- fe means field element.  Here the field is \Z/(2^255-19).
-- An element t, entries t[0]...t[9], represents the integer
-- t[0]+2^26 t[1]+2^51 t[2]+2^77 t[3]+2^102 t[4]+...+2^230 t[9].
-- Bounds on each t[i] vary depending on context.

-- mont_pub_to_ed_pub
toSigningKey :: X25519.PublicKey -> PublicKey
toSigningKey mont_pub0 = PublicKey ed_pub
 where
    -- Read the public key as a field element
    mont_pub = fe_frombytes mont_pub0

    -- Convert the Montgomery public key to a twisted Edwards public key
    fe_ONE = fe_1

    -- Calculate the parameters (u - 1) and (u + 1)
    mont_pub_minus_one = fe_sub mont_pub fe_ONE
    mont_pub_plus_one0 = fe_add mont_pub fe_ONE

    -- Prepare inv(u + 1)
    mont_pub_plus_one = fe_invert mont_pub_plus_one0

    -- Calculate y = (u - 1) * inv(u + 1) (mod p)
    ed_pub0 = fe_mul mont_pub_minus_one mont_pub_plus_one
    ed_pub = fe_tobytes ed_pub0

-- mont_priv_to_ed_pair
toSigningKeyPair :: X25519.SecretKey -> (SecretKey,PublicKey)
toSigningKeyPair mont_priv0 = (SecretKey ed_priv, PublicKey ed_pub)
 where
    -- Prepare a buffer for the twisted Edwards private key
    ed_priv1 = (throwCryptoError . scalarDecodeLong :: X25519.SecretKey -> Scalar) mont_priv0

    -- Get the twisted edwards public key, including the sign bit
    ed_pub0 = ge_p3_tobytes $ ge_scalarmult_base ed_priv1

    -- Save the sign bit for later
    sign_bit = (BA.index ed_pub0 31 `shiftR` 7) .&. 1

    -- Force the sign bit to zero
    pub' :: Bytes
    pub' = BA.copyAndFreeze ed_pub0 $ \ptr -> do
        let at31 = plusPtr ptr 31
        byte31 <- peek at31
        poke at31 $ (byte31 .&. 0x7F) `asTypeOf` sign_bit

    CryptoPassed ed_pub = Ed25519.publicKey pub'


    -- Prepare the negated private key
    ed_priv_neg = sc_neg ed_priv1

    -- Get the correct private key based on the sign stored above
    ed_priv = if sign_bit/=0 then ed_priv_neg
                             else ed_priv1

-- sc_zero = throwCryptoError $ scalarDecodeLong (b::Bytes)
--  where
--     b = BA.pack $ encodeLittleEndian $ 2^252 + 27742317777372353535851937790883648493

sc_neg1 :: Scalar
sc_neg1 = throwCryptoError $ scalarDecodeLong (b::Bytes)
 where
    b = BA.pack $ encodeLittleEndian $ 2^252 + 27742317777372353535851937790883648492