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

module Eth.AddressSpec (spec) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Base16 as B16
import Eth.Address (verifySignature)
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 "verifySignature" $ do
    describe "input validation" $ do
      it "rejects signature with wrong length" $ do
        let message = "test message"
            sig = BS.replicate 64 0 -- Should be 65
            addr = BS.replicate 20 0
        result <- verifySignature message sig addr
        result `shouldBe` False

      it "rejects address with wrong length" $ do
        let message = "test message"
            sig = BS.replicate 65 0
            addr = BS.replicate 19 0 -- Should be 20
        result <- verifySignature message sig addr
        result `shouldBe` False

      it "accepts correct lengths (even if signature is invalid)" $ do
        let message = "test message"
            sig = BS.replicate 65 0
            addr = BS.replicate 20 0
        -- Should not crash, just return False for invalid signature
        result <- verifySignature message sig addr
        result `shouldBe` False

    describe "property-based tests" $ do
      it "always returns a boolean" $
        property $
          \(message :: BS.ByteString) (sig :: BS.ByteString) (addr :: BS.ByteString) -> monadicIO $ do
            result <- run $ verifySignature message sig addr
            return $ result `elem` [True, False]

      it "returns False for signatures with wrong length" $
        property $
          forAll (choose (0, 128) `suchThat` (/= 65)) $ \sigLen -> monadicIO $ do
            let message = "test"
                sig = BS.replicate sigLen 0
                addr = BS.replicate 20 0
            result <- run $ verifySignature message sig addr
            return $ result == False

      it "returns False for addresses with wrong length" $
        property $
          forAll (choose (0, 40) `suchThat` (/= 20)) $ \addrLen -> monadicIO $ do
            let message = "test"
                sig = BS.replicate 65 0
                addr = BS.replicate addrLen 0
            result <- run $ verifySignature message sig addr
            return $ result == False

      it "is deterministic" $
        property $
          \(message :: BS.ByteString) -> monadicIO $ do
            let sig = BS.replicate 65 0
                addr = BS.replicate 20 0
            result1 <- run $ verifySignature message sig addr
            result2 <- run $ verifySignature message sig addr
            return $ result1 == result2

    describe "signature verification" $ do
      it "rejects an invalid signature for correct message/address" $ do
        let message = "Hello, Web3!"
            sig' = B16.decode "9d2db20f926fad885854f4bf659a72de9db2990b199ad300644e2b482f8583e169325c2a57bbfa4ef6897f63431f290716c58507af64d88554923410f38be2e31b"
            addr' = B16.decode "f406ffc8e58c5163e67e020dd1c5ce924923ecbf"
        result <- case (sig', addr') of
          (Right sig, Right addr) -> verifySignature message sig addr
          (_, _) -> error "Couldn't decode signature or address!"
        result `shouldBe` False

      it "rejects invalid message from with valid address/signature" $ do
        let message = "Bye, Web3!"
            sig' = B16.decode "8d2db20f926fad885854f4bf659a72de9db2990b199ad300644e2b482f8583e169325c2a57bbfa4ef6897f63431f290716c58507af64d88554923410f38be2e31b"
            addr' = B16.decode "f406ffc8e58c5163e67e020dd1c5ce924923ecbf"
        result <- case (sig', addr') of
          (Right sig, Right addr) -> verifySignature message sig addr
          (_, _) -> error "Couldn't decode signature or address!"
        result `shouldBe` False

      it "rejects invalid address from with valid message/signature [from another address]" $ do
        let message = "Hello, Web3!"
            sig' = B16.decode "8d2db20f926fad885854f4bf659a72de9db2990b199ad300644e2b482f8583e169325c2a57bbfa4ef6897f63431f290716c58507af64d88554923410f38be2e31b"
            addr' = B16.decode "0406ffc8e58c5163e67e020dd1c5ce924923ecbf"
        result <- case (sig', addr') of
          (Right sig, Right addr) -> verifySignature message sig addr
          (_, _) -> error "Couldn't decode signature or address!"
        result `shouldBe` False

      it "verifies a valid Ethereum signature from real wallet" $ do
        let message = "Hello, Web3!"
            sig' = B16.decode "8d2db20f926fad885854f4bf659a72de9db2990b199ad300644e2b482f8583e169325c2a57bbfa4ef6897f63431f290716c58507af64d88554923410f38be2e31b"
            addr' = B16.decode "f406ffc8e58c5163e67e020dd1c5ce924923ecbf"
        result <- case (sig', addr') of
          (Right sig, Right addr) -> verifySignature message sig addr
          (_, _) -> error "Couldn't decode signature or address!"
        result `shouldBe` True

    describe "edge cases" $ do
      it "handles empty message" $ do
        let message = ""
            sig = BS.replicate 65 0
            addr = BS.replicate 20 0
        result <- verifySignature message sig addr
        result `shouldSatisfy` const True -- Just shouldn't crash
      it "handles very long message" $ do
        let message = BS.replicate 10000 0xff
            sig = BS.replicate 65 0
            addr = BS.replicate 20 0
        result <- verifySignature message sig addr
        result `shouldSatisfy` const True -- Just shouldn't crash