{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Ecc where

import qualified Crypto.Hash.SHA256 as SHA256
import           Data.Bits
import qualified Data.ByteString    as BS
import           Data.Proxy
import           GHC.TypeLits
import           Text.Printf        (PrintfArg, printf)

-- FiniteFields
--https://stackoverflow.com/questions/39823408/prime-finite-field-z-pz-in-haskell-with-operator-overloading
newtype FieldElement (n :: Nat) = FieldElement Integer deriving Eq

instance KnownNat n => Num (FieldElement n) where
  FieldElement x + FieldElement y = fromInteger $ x + y
  FieldElement x * FieldElement y = fromInteger $ x * y
  abs x = x
  signum _ = 1
  negate (FieldElement x) = fromInteger $ negate x
  fromInteger a = FieldElement (mod a n) where n = natVal (Proxy :: Proxy n)

instance KnownNat n => Fractional (FieldElement n) where
  recip a = a ^ (n - 2) where n = natVal (Proxy :: Proxy n)
  fromRational r = error "can't transform" -- fromInteger (numerator r) / fromInteger (denominator r)

instance KnownNat n => Show (FieldElement n) where
  show (FieldElement a) | n == (2 ^ 256 - 2 ^ 32 - 977) = printf "0x%064x" a
                        | otherwise = "FieldElement_" ++ show n ++ " " ++ show a
    where n = natVal (Proxy :: Proxy n)

-- Elliptic curve
data ECPoint a
  = Infinity
  | ECPoint
      { x :: a
      , y :: a
      , a :: a
      , b :: a
      }
  deriving (Eq)

instance {-# OVERLAPPABLE #-} (PrintfArg a, Num a) => Show (ECPoint a) where
  show Infinity = "ECPoint(Infinity)"
  show p        = printf "ECPoint(%f, %f)_%f_%f" (x p) (y p) (a p) (b p)

instance {-# OVERLAPPING  #-} KnownNat n => Show (ECPoint (FieldElement n)) where
  show Infinity = "ECPoint(Infinity)"
  show p | n == (2 ^ 256 - 2 ^ 32 - 977) = "S256Point" ++ points
         | otherwise = "ECPoint_" ++ show n ++ points ++ params
   where
    n      = natVal (Proxy :: Proxy n)
    points = "(" ++ si (x p) ++ ", " ++ si (y p) ++ ")"
    params = "a_" ++ si (a p) ++ "|b_" ++ si (b p)
    si (FieldElement r) | n == (2 ^ 256 - 2 ^ 32 - 977) = printf "0x%064x" r
                        | otherwise                     = show r

validECPoint :: (Eq a, Num a) => ECPoint a -> Bool
validECPoint Infinity          = True
validECPoint (ECPoint x y a b) = y ^ 2 == x ^ 3 + a * x + b

add :: (Eq a, Fractional a) => ECPoint a -> ECPoint a -> ECPoint a
add Infinity p        = p
add p        Infinity = p
add p q | a p /= a q || b p /= b q = error "point not on same curve"
        | x p == x q && y p /= y q = Infinity
        | x p /= x q               = new_point $ (y q - y p) / (x q - x p)
        | x p == x q && y p == 0   = Infinity
        | p == q                   = new_point $ (3 * x p ^ 2 + a p) / (2 * y p)
        | otherwise                = error "Unexpected case of points"
 where
  new_point slope =
    let new_x = slope ^ 2 - x p - x q
        new_y = slope * (x p - new_x) - y p
    in  ECPoint new_x new_y (a p) (b p)


binaryExpansion :: (Semigroup a) => Integer -> a -> a -> a
binaryExpansion m value result
  | m == 0    = result
  | otherwise = binaryExpansion (m `shiftR` 1) (value <> value) accumulator
  where accumulator = if m .&. 1 == 1 then result <> value else result

scalarProduct :: (Eq a, Fractional a) => Integer -> ECPoint a -> ECPoint a
scalarProduct m ec = binaryExpansion m ec Infinity

instance (Eq a, Fractional a) => Semigroup (ECPoint a) where
  (<>) = add

instance (Eq a, Fractional a) => Monoid (ECPoint a) where
  mempty = Infinity

-- The Bitcoin curve
type S256Field = FieldElement (2 ^ 256- 2^ 32 - 977)
type NField
  = FieldElement
      0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141
type S256Point = ECPoint S256Field

s256point :: S256Field -> S256Field -> S256Point
s256point x y =
  let p = ECPoint x y 0 7
  in  if validECPoint p then p else error "Invalid point"

ncons = 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141
gcons = s256point
  0x79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798
  0x483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8

-- ECDSA
data Signature = Signature
  { r :: S256Field
  , s :: NField
  }
  deriving Show

asInt :: KnownNat n => FieldElement n -> Integer
asInt (FieldElement n) = n

verifySignanture :: NField -> Signature -> S256Point -> Bool
verifySignanture z (Signature r s) pub = x target == r
 where
  target =
    scalarProduct (asInt $ z / s) gcons
      <> scalarProduct (asInt $ fromInteger (asInt r) / s) pub

hash256 :: BS.ByteString -> BS.ByteString
hash256 = SHA256.hash . SHA256.hash

fromBytes :: BS.ByteString -> Integer
fromBytes = BS.foldl' f 0 where f a b = a `shiftL` 8 .|. fromIntegral b

sighash :: BS.ByteString -> NField
sighash = fromInteger . fromBytes . hash256

integerToBytes :: Integer -> BS.ByteString
integerToBytes = BS.pack . go
 where
  go c = case c of
    0 -> []
    c -> go (c `div` 256) ++ [fromIntegral (c `mod` 256)]

zeroPad :: Integer -> BS.ByteString -> BS.ByteString
zeroPad n s = BS.append padding s
 where
  padding = BS.pack (replicate (fromIntegral n - fromIntegral (BS.length s)) 0)

toBytes32 :: Integer -> BS.ByteString
toBytes32 = zeroPad 32 . integerToBytes

deterministicK :: NField -> NField -> NField
deterministicK (FieldElement priv) (FieldElement z) = fromInteger
  $ candidate k2 v2
 where
  k    = BS.pack $ replicate 32 0
  v    = BS.pack $ replicate 32 1
  zbs  = toBytes32 z
  skbs = toBytes32 priv
  k1   = SHA256.hmac k $ v `BS.append` "\NUL" `BS.append` skbs `BS.append` zbs
  v1   = SHA256.hmac k1 v
  k2   = SHA256.hmac k1 $ v1 `BS.append` "\SOH" `BS.append` skbs `BS.append` zbs
  v2   = SHA256.hmac k2 v1

  candidate k v =
    let vNew = SHA256.hmac k v
        can  = fromBytes vNew
    in  if can >= 1 && can < ncons
          then can
          else
            let kp = SHA256.hmac k $ vNew `BS.append` "\NUL"
                vp = SHA256.hmac kp vNew
            in  candidate kp vp

signMessage :: NField -> BS.ByteString -> Signature
signMessage priv mesg =
  let z               = sighash mesg
      k               = deterministicK priv z
      rm              = scalarProduct (asInt k) gcons
      FieldElement sm = (z + fromInteger (asInt (x rm)) * priv) / k
      ss              = if sm > div ncons 2 then ncons - sm else sm
  in  Signature (x rm) (fromInteger ss)


-- TESTs
-- Finite Fields
assert :: Bool -> Bool
assert False = error "WRONG"
assert x     = x

aa =
  let a = FieldElement 2 :: FieldElement 31
      b = FieldElement 15
  in  (a + b == FieldElement 17, a /= b, a - b == FieldElement 18)

bb =
  let a = FieldElement 19 :: FieldElement 31
      b = FieldElement 24
  in  a * b
-- Elliptic curves
tre = FieldElement 3 :: FieldElement 31
cc =
  let a = ECPoint tre (-7) 5 7
      b = ECPoint 18 77 5 7
      c = ECPoint (-1) (-1) 5 7
  in  ( validECPoint a
      , validECPoint b
      , validECPoint c
      , a /= b
      , a == a
      , add Infinity          a
      , add a                 (ECPoint 3 7 5 7)
      , add (ECPoint 3 7 5 7) c
      , add c                 c
      )

dd =
  let a = FieldElement 0 :: FieldElement 223
      b = FieldElement 7
      x = FieldElement 192
      y = FieldElement 105
  in  ECPoint x y a b
ee = ECPoint 192 105 (FieldElement 0 :: FieldElement 223) 7
ff = ECPoint 192 105 0 7 :: ECPoint (FieldElement 223)

aPoint = ECPoint 192 105 0 7 :: ECPoint (FieldElement 223)
total = add aPoint $ add aPoint $ add aPoint $ add aPoint aPoint

totalfold = foldr add Infinity $ replicate 5 aPoint
totalmconcat = mconcat $ replicate 5 aPoint

-- Bitcoin curve
li :: S256Field
li = 12
ll :: ECPoint (FieldElement 31)
ll = Infinity
ri = ECPoint 3 7 5 7 :: S256Point

-- ECDSA

-- z = 0xbc62d4b80d9e36da29c16c5d4d9f11731f36052c72401a76c23c0fb5a9b74423
-- r = 0x37206a0610995c58074999cb9767b87af4c4978db68c06e8e6e81d282047a7c6
-- s = 0x8ca63759c1157ebeaec0d03cecca119fc9a75bf8e6d0fa65c841c8e2738cdaec ::NField
-- px = 0x04519fac3d910ca7e7138f7013706f619fa8f033e6ec6e09370ea38cee6a7574
-- py = 0x82b51eab8c27c66e26c858a079bcdf4f1ada34cec420cafc7eac1a42216fb6c4
-- point = s256point px py
-- u = z / s
-- v = r / s
-- signa = scalarProduct (asInt u) gcons <> scalarProduct (asInt v) point

pub = s256point
  0x887387e452b8eacc4acfde10d9aaf7f6d9a0f975aabb10d006e4da568744d06c
  0x61de6d95231cd89026e286df3b6ae4a894a3378e393e93a0f45b666329a0ae34

z1 =
  0xec208baa0fc1c19f708a9ca96fdeff3ac3f230bb4a7ba4aede4942ad003c0f60 :: NField
r1 =
  0xac8d1c87e51d0d441be8b3dd5b05c8795b48875dffe00b7ffcfac23010d3a395 :: S256Field
s1 =
  0x68342ceff8935ededd102dd876ffd6ba72d6a427a3edb13d26eb0781cb423c4 :: NField

--signa1 =
  --scalarProduct (asInt $ z1 / s1) gcons <> scalarProduct (asInt $ r1 / s1) pub

z2 =
  0x7c076ff316692a3d7eb3c3bb0f8b1488cf72e1afcd929e29307032997a838a3d :: NField
r2 = 0xeff69ef2b1bd93a66ed5219add4fb51e11a840f404876325a1e8ffe0529a2c :: NField
s2 =
  0xc7207fee197d27c618aea621406f6bf5ef6fca38681d82b2f06fddbdce6feab6 :: NField

-- priv = fromIntegral $ fromBytes $ hash256 "my secret" :: NField
priv = 12345
mesg = sighash "Programming Bitcoin!"
k = 1234567890 :: NField
rm = scalarProduct (asInt k) gcons
sm = (mesg + fromIntegral (asInt (x rm)) * priv) / k
pubm = scalarProduct (asInt priv) gcons