{-# LANGUAGE OverloadedStrings #-}

module UpgradeSpec (spec) where

import Arbitrary ()
import Control.Applicative (Alternative ((<|>)))
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import qualified Data.Set as Set
import OpenTimestamps.Attestation (Attestation (Unknown))
import OpenTimestamps.Op ()
import OpenTimestamps.Timestamp (Timestamp (..), merge)
import OpenTimestamps.Upgrade (mergeIntoSubTimestamp)
import Test.Hspec (Spec, describe, it, shouldBe)
import Test.QuickCheck
  ( Arbitrary (arbitrary)
  , Gen
  , Testable (property)
  , discard
  , elements
  , forAll
  , suchThat
  )

getAllMessages :: Timestamp -> Set.Set BS.ByteString
getAllMessages ts = Set.insert (timestampMsg ts) (foldMap getAllMessages (Map.elems (ops ts)))

genTargetMsgInTimestamp :: Timestamp -> Gen BS.ByteString
genTargetMsgInTimestamp ts = do
  let allMsgs = getAllMessages ts
  if Set.null allMsgs
    then arbitrary -- Should not happen with valid Timestamp generation, but as a fallback
    else elements (Set.toList allMsgs)

findTimestamp :: BS.ByteString -> Timestamp -> Maybe Timestamp
findTimestamp target ts
  | timestampMsg ts == target = Just ts
  | otherwise =
      foldr
        ( \(_op, subTs) acc ->
            findTimestamp target subTs <|> acc
        )
        Nothing
        (Map.toList (ops ts))

spec :: Spec
spec = do
  describe "mergeIntoSubTimestamp" $ do
    it "returns True if targetMsg is found and merged with new information" $ property $ \currentTs ->
      forAll (genTargetMsgInTimestamp currentTs) $ \targetMsg ->
        forAll arbitrary $ \newAtt ->
          -- Ensure newAtt is not already in the sub-timestamp at targetMsg
          let subTs = fromMaybe currentTs (findTimestamp targetMsg currentTs) -- Fallback
              uniqueAtt = if Set.member newAtt (attestations subTs) then Unknown "" "" else newAtt
              newTs = Timestamp targetMsg (Set.singleton uniqueAtt) Map.empty
              (_, changed) = mergeIntoSubTimestamp targetMsg currentTs newTs
           in changed `shouldBe` True

    it "returns False if targetMsg is not found" $ property $ \currentTs newTs ->
      forAll (suchThat arbitrary (\msg -> not (Set.member msg (getAllMessages currentTs)))) $ \targetMsg ->
        let (_, changed) = mergeIntoSubTimestamp targetMsg currentTs newTs
         in changed `shouldBe` False

    it "merges correctly when targetMsg matches root" $ property $ \currentTs newTs ->
      let targetMsg = timestampMsg currentTs
          (mergedTs, _) = mergeIntoSubTimestamp targetMsg currentTs newTs
       in do
            mergedTs `shouldBe` fst (merge currentTs newTs)
            pure ()

    it "preserves structure for unaffected parts" $ property $ \currentTs newTs ->
      forAll (genTargetMsgInTimestamp currentTs) $ \targetMsg ->
        if targetMsg == timestampMsg currentTs
          then discard -- Covered by "merges correctly when targetMsg matches root"
          else
            let (mergedTs, _) = mergeIntoSubTimestamp targetMsg currentTs newTs
             in do
                  timestampMsg mergedTs `shouldBe` timestampMsg currentTs
                  attestations mergedTs `shouldBe` attestations currentTs

{- TODO This fails intermittently, disable for now & investigate later
it "returns False for changed flag if newTs adds no new information at target" $ property $ \currentTs ->
  forAll (genTargetMsgInTimestamp currentTs) $ \targetMsg ->
    -- Ensure targetMsg is unique in currentTs for this test
    let allMsgs = getAllMessages currentTs
        countTargetMsg = length $ filter (== targetMsg) (Set.toList allMsgs)
     in if countTargetMsg > 1
          then discard
          else case findTimestamp targetMsg currentTs of
            Just subTsToMerge -> do
              -- Verify `merge` itself is idempotent for identical timestamps
              let (mergedSelf, selfChanged) = merge subTsToMerge subTsToMerge
              selfChanged `shouldBe` False
              mergedSelf `shouldBe` subTsToMerge

              -- Now, for the overall mergeIntoSubTimestamp:
              -- If the sub-timestamp at targetMsg is merged with an identical one,
              -- the overall 'changed' flag should be False.
              let (_, overallChanged) = mergeIntoSubTimestamp targetMsg currentTs subTsToMerge
              overallChanged `shouldBe` False
            Nothing -> discard -- Should not happen with genTargetMsgInTimestamp
-}