{-# LANGUAGE OverloadedStrings #-}

module SignVerifyTweakProperty (propertySignVerifyTweak) where

import Crypto.Curve.Secp256k1 (derive_pub)
import Crypto.Curve.Secp256k1.MuSig2 (SecKey (..), SecNonce (..), Tweak (..), aggNonces, mkSessionContext, partialSigVerify, publicNonce, sign)
import Crypto.Curve.Secp256k1.MuSig2.Internal (curveOrder)
import Data.ByteString (ByteString)
import Data.Maybe (fromJust, fromMaybe)
import Test.Tasty
import Test.Tasty.QuickCheck
import Util ()

propertySignVerifyTweak :: TestTree
propertySignVerifyTweak =
  testGroup
    "sign and partialSigVerify with Tweaks Properties"
    [ testProperty "Valid Signature Range with Tweaks" prop_validSignatureRangeWithTweaks
    , testProperty "Sign-Verify Roundtrip with Tweaks" prop_signVerifyRoundtripWithTweaks
    , testProperty "Signature Determinism with Tweaks" prop_signatureDeterminismWithTweaks
    , testProperty "Empty Tweaks Equals No Tweaks" prop_emptyTweaksEqualsNoTweaks
    , testProperty "Plain Tweaks Are Commutative" prop_plainTweaksCommutative
    ]

-- | Property: Generated signatures with tweaks are in the valid range \([0, Q-1]\).
prop_validSignatureRangeWithTweaks :: SecNonce -> SecKey -> Property
prop_validSignatureRangeWithTweaks secNonce secKey@(SecKey sk) =
  forAll (resize 5 $ listOf arbitrary) $ \tweaks ->
    forAll arbitrary $ \msg ->
      let pubkey = fromMaybe (error "Failed to derive pubkey") $ derive_pub (fromInteger sk)
          pubNonce = publicNonce secNonce
          pubNonces = [pubNonce]
          pubkeys = [pubkey]
          aggNonce = fromJust $ aggNonces pubNonces
          ctx = mkSessionContext aggNonce pubkeys tweaks msg
          sig = sign secNonce secKey ctx
       in (sig >= 0) .&&. (sig < curveOrder)

-- | Property: A signature created with tweaks verifies with 'partialSigVerify'.
prop_signVerifyRoundtripWithTweaks :: SecNonce -> SecKey -> Property
prop_signVerifyRoundtripWithTweaks secNonce secKey@(SecKey sk) =
  forAll (resize 5 $ listOf arbitrary) $ \tweaks ->
    forAll arbitrary $ \msg ->
      let pubkey = fromMaybe (error "Failed to derive pubkey") $ derive_pub (fromInteger sk)
          pubNonce = publicNonce secNonce
          pubNonces = [pubNonce]
          pubkeys = [pubkey]
          aggNonce = fromJust $ aggNonces pubNonces
          ctx = mkSessionContext aggNonce pubkeys tweaks msg
          sig = sign secNonce secKey ctx
          signerIndex = 0 -- We're always the first signer in this test
          result = partialSigVerify sig pubNonces pubkeys tweaks msg signerIndex
       in result === True

-- | Property: Signing the same message with the same tweaks produces the same signature.
prop_signatureDeterminismWithTweaks :: SecNonce -> SecKey -> Property
prop_signatureDeterminismWithTweaks secNonce secKey@(SecKey sk) =
  forAll (resize 5 $ listOf arbitrary) $ \tweaks ->
    forAll arbitrary $ \msg ->
      let pubkey = fromMaybe (error "Failed to derive pubkey") $ derive_pub (fromInteger sk)
          pubNonce = publicNonce secNonce
          pubNonces = [pubNonce]
          pubkeys = [pubkey]
          aggNonce = fromJust $ aggNonces pubNonces
          ctx = mkSessionContext aggNonce pubkeys tweaks msg
          sig1 = sign secNonce secKey ctx
          sig2 = sign secNonce secKey ctx
       in sig1 === sig2

-- | Property: Empty tweaks should produce the same result as no tweaks.
prop_emptyTweaksEqualsNoTweaks :: SecNonce -> SecKey -> ByteString -> Property
prop_emptyTweaksEqualsNoTweaks secNonce secKey@(SecKey sk) msg =
  let pubkey = fromMaybe (error "Failed to derive pubkey") $ derive_pub (fromInteger sk)
      pubNonce = publicNonce secNonce
      pubNonces = [pubNonce]
      pubkeys = [pubkey]
      aggNonce = fromJust $ aggNonces pubNonces

      -- Context with empty tweaks
      ctxWithEmptyTweaks = mkSessionContext aggNonce pubkeys [] msg
      sigWithEmptyTweaks = sign secNonce secKey ctxWithEmptyTweaks

      -- Context with no tweaks parameter (though we still pass empty list)
      ctxWithNoTweaks = mkSessionContext aggNonce pubkeys [] msg
      sigWithNoTweaks = sign secNonce secKey ctxWithNoTweaks
   in sigWithEmptyTweaks === sigWithNoTweaks

-- | Property: Plain tweaks are commutative.
prop_plainTweaksCommutative :: SecNonce -> SecKey -> Integer -> Integer -> ByteString -> Property
prop_plainTweaksCommutative secNonce secKey@(SecKey sk) t1 t2 msg =
  t1 > 0
    && t1 < curveOrder
    && t2 > 0
    && t2 < curveOrder
    && t1
      /= t2
    ==> let pubkey = fromMaybe (error "Failed to derive pubkey") $ derive_pub (fromInteger sk)
            pubNonce = publicNonce secNonce
            pubNonces = [pubNonce]
            pubkeys = [pubkey]
            aggNonce = fromJust $ aggNonces pubNonces

            -- First order: [PlainTweak t1, PlainTweak t2]
            tweak1 = PlainTweak t1
            tweak2 = PlainTweak t2
            ctx1 = mkSessionContext aggNonce pubkeys [tweak1, tweak2] msg
            sig1 = sign secNonce secKey ctx1

            -- Second order: [PlainTweak t2, PlainTweak t1]
            ctx2 = mkSessionContext aggNonce pubkeys [tweak2, tweak1] msg
            sig2 = sign secNonce secKey ctx2
         in -- Plain tweaks are commutative, order shouldn't matter
            sig1 === sig2