{-# LANGUAGE OverloadedStrings #-}

module PruneSpec (spec) where

import Arbitrary ()
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import OpenTimestamps.Attestation (Attestation (Bitcoin, Pending))
import OpenTimestamps.Op (Op (Append, Prepend, Sha256))
import OpenTimestamps.Prune (pruneTimestamp)
import OpenTimestamps.Timestamp (Timestamp (..), getAttestations, isTimestampComplete)
import Test.Hspec
import Test.QuickCheck

spec :: Spec
spec = describe "OpenTimestamps.Prune" $ do
  it "prunes a simple timestamp with redundant Bitcoin attestations" $ do
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 100, Bitcoin 200, Pending "http://example.com"]
            , ops = Map.empty
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 100, Pending "http://example.com"]
            , ops = Map.empty
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  it "prunes a timestamp with redundant Bitcoin attestations in sub-timestamps" $ do
    let subTs1 =
          Timestamp
            { timestampMsg = "subMsg1"
            , attestations = Set.fromList [Bitcoin 300, Pending "http://calendar1.org"]
            , ops = Map.empty
            }
    let subTs2 =
          Timestamp
            { timestampMsg = "subMsg2"
            , attestations = Set.fromList [Bitcoin 100, Pending "http://calendar2.org"]
            , ops = Map.empty
            }
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 200]
            , ops = Map.fromList [(Sha256, subTs1), (Append "data", subTs2)]
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [] -- Bitcoin 200 is suboptimal, removed
            , ops =
                Map.fromList
                  [
                    ( Sha256
                    , Timestamp
                        { timestampMsg = "subMsg1"
                        , attestations = Set.fromList [Pending "http://calendar1.org"]
                        , ops = Map.empty
                        }
                    )
                  ,
                    ( Append "data"
                    , Timestamp
                        { timestampMsg = "subMsg2"
                        , attestations = Set.fromList [Bitcoin 100, Pending "http://calendar2.org"]
                        , ops = Map.empty
                        }
                    )
                  ]
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  it "prunes a timestamp with empty branches" $ do
    let emptyTs =
          Timestamp
            { timestampMsg = "emptyMsg"
            , attestations = Set.empty
            , ops = Map.empty
            }
    let subTs =
          Timestamp
            { timestampMsg = "subMsg"
            , attestations = Set.fromList [Pending "http://calendar.org"]
            , ops = Map.singleton Sha256 emptyTs
            }
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.empty
            , ops = Map.singleton (Append "data") subTs
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.empty
            , ops =
                Map.singleton
                  (Append "data")
                  ( Timestamp
                      { timestampMsg = "subMsg"
                      , attestations = Set.fromList [Pending "http://calendar.org"]
                      , ops = Map.empty -- EmptyTs should be pruned
                      }
                  )
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  it "prunes a complex timestamp with multiple redundant Bitcoin attestations and empty branches" $ do
    let subSubTs1 =
          Timestamp
            { timestampMsg = "subSubMsg1"
            , attestations = Set.fromList [Bitcoin 500]
            , ops = Map.empty
            }
    let subSubTs2 =
          Timestamp
            { timestampMsg = "subSubMsg2"
            , attestations = Set.fromList [Bitcoin 50]
            , ops = Map.empty
            }
    let subTs1 =
          Timestamp
            { timestampMsg = "subMsg1"
            , attestations = Set.fromList [Bitcoin 400, Pending "http://cal1.org"]
            , ops = Map.singleton Sha256 subSubTs1
            }
    let subTs2 =
          Timestamp
            { timestampMsg = "subMsg2"
            , attestations = Set.fromList [Pending "http://cal2.org"]
            , ops = Map.singleton (Prepend "prefix") subSubTs2
            }
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 1000, Pending "http://cal0.org"]
            , ops = Map.fromList [(Append "data1", subTs1), (Append "data2", subTs2)]
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Pending "http://cal0.org"]
            , ops =
                Map.fromList
                  [
                    ( Append "data1"
                    , Timestamp
                        { timestampMsg = "subMsg1"
                        , attestations = Set.fromList [Pending "http://cal1.org"]
                        , ops = Map.empty -- Sha256 op and subSubTs1 removed
                        }
                    )
                  ,
                    ( Append "data2"
                    , Timestamp
                        { timestampMsg = "subMsg2"
                        , attestations = Set.fromList [Pending "http://cal2.org"]
                        , ops =
                            Map.fromList
                              [
                                ( Prepend "prefix"
                                , Timestamp
                                    { timestampMsg = "subSubMsg2"
                                    , attestations = Set.fromList [Bitcoin 50]
                                    , ops = Map.empty
                                    }
                                )
                              ]
                        }
                    )
                  ]
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  it "prunes a timestamp with Bitcoin attestations of same height but different depths" $ do
    let subTs =
          Timestamp
            { timestampMsg = "subMsg"
            , attestations = Set.fromList [Bitcoin 100, Pending "http://calendar.org"]
            , ops = Map.empty
            }
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 100]
            , ops = Map.singleton Sha256 subTs
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Bitcoin 100] -- This Bitcoin 100 is at depth 0, better than subMsg's Bitcoin 100 at depth 1
            , ops =
                Map.singleton
                  Sha256
                  ( Timestamp
                      { timestampMsg = "subMsg"
                      , attestations = Set.fromList [Pending "http://calendar.org"]
                      , ops = Map.empty
                      }
                  )
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  it "prunes a timestamp with no Bitcoin attestations" $ do
    let subTs =
          Timestamp
            { timestampMsg = "subMsg"
            , attestations = Set.fromList [Pending "http://calendar.org"]
            , ops = Map.empty
            }
    let initialTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Pending "http://example.com"]
            , ops = Map.singleton Sha256 subTs
            }
    let expectedTs =
          Timestamp
            { timestampMsg = "msg1"
            , attestations = Set.fromList [Pending "http://example.com"]
            , ops = Map.singleton Sha256 subTs
            }
    pruneTimestamp initialTs `shouldBe` expectedTs

  describe "pruneTimestamp properties" $ do
    it "preserves completeness" $
      quickCheck propPrunePreservesCompleteness
    it "reduces size appropriately" $
      quickCheck propPruneReducesSize
    it "is idempotent" $
      quickCheck propPruneIdempotent

-- | Pruning should preserve timestamp completeness
propPrunePreservesCompleteness :: Timestamp -> Bool
propPrunePreservesCompleteness ts =
  isTimestampComplete (pruneTimestamp ts) == isTimestampComplete ts

-- | Pruning should reduce or maintain the number of attestations
propPruneReducesSize :: Timestamp -> Bool
propPruneReducesSize ts =
  let originalAtts = getAttestations ts
      prunedAtts = getAttestations (pruneTimestamp ts)
   in length prunedAtts <= length originalAtts

-- | Pruning should be idempotent
propPruneIdempotent :: Timestamp -> Bool
propPruneIdempotent ts =
  pruneTimestamp (pruneTimestamp ts) == pruneTimestamp ts