{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Crypto.Secp256k1.RecoverySpec (spec) where

import Crypto.Secp256k1.Recovery (recoverPubKey)
import qualified Data.ByteString as BS
import Data.Maybe (isJust, isNothing)
import Test.Hspec
import Test.QuickCheck
import Test.QuickCheck.Monadic

-- Arbitrary instance for ByteString
instance Arbitrary BS.ByteString where
  arbitrary = BS.pack <$> arbitrary

spec :: Spec
spec = do
  describe "recoverPubKey" $ do
    describe "input validation" $ do
      it "rejects signature that is too short" $ do
        let message = BS.replicate 32 0
            sig = BS.replicate 64 0 -- Should be 65
        result <- recoverPubKey message sig
        result `shouldBe` Nothing

      it "rejects signature that is too long" $ do
        let message = BS.replicate 32 0
            sig = BS.replicate 66 0 -- Should be 65
        result <- recoverPubKey message sig
        result `shouldBe` Nothing

      it "rejects message that is too short" $ do
        let message = BS.replicate 31 0 -- Should be 32
            sig = BS.replicate 65 0
        result <- recoverPubKey message sig
        result `shouldBe` Nothing

      it "rejects message that is too long" $ do
        let message = BS.replicate 33 0 -- Should be 32
            sig = BS.replicate 65 0
        result <- recoverPubKey message sig
        result `shouldBe` Nothing

    describe "property-based tests" $ do
      it "always returns Nothing or Just a 65-byte public key" $ property $
        \(message :: BS.ByteString) (sig :: BS.ByteString) -> monadicIO $ do
          result <- run $ recoverPubKey message sig
          monitor $ classify (isNothing result) "returned Nothing"
          monitor $ classify (isJust result) "returned Just pubKey"
          case result of
            Nothing -> return True
            Just pubKey -> return $ BS.length pubKey == 65

      it "returns Nothing for invalid signature lengths" $ property $
        forAll (choose (0, 128) `suchThat` (/= 65)) $ \len -> monadicIO $ do
          let message = BS.replicate 32 0
              sig = BS.replicate len 0
          result <- run $ recoverPubKey message sig
          return $ result == Nothing

      it "returns Nothing for invalid message lengths" $ property $
        forAll (choose (0, 64) `suchThat` (/= 32)) $ \len -> monadicIO $ do
          let message = BS.replicate len 0
              sig = BS.replicate 65 0
          result <- run $ recoverPubKey message sig
          return $ result == Nothing