-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Tests for operation size evaluation.

module Test.OpSize
  ( test_Nesting
  , test_Numbers
  , test_Values
  , test_Types
  , test_Instructions
  ) where

import Prelude hiding (EQ)

import Data.Text qualified as T
import Fmt (pretty)
import Hedgehog (Gen)
import Test.HUnit (Assertion, (@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase)

import Hedgehog.Gen.Tezos.Address (genAddress)
import Hedgehog.Gen.Tezos.Crypto (genKeyHash)
import Hedgehog.Gen.Tezos.Crypto.Ed25519 qualified as Ed25519
import Hedgehog.Gen.Tezos.Crypto.Secp256k1 qualified as Secp256k1
import Morley.Michelson.Macro
import Morley.Michelson.Parser qualified as Parser
import Morley.Michelson.Typed
import Morley.Michelson.Untyped (buildEpName)
import Morley.Michelson.Untyped.OpSize
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Util.Text (dquotes)
import Test.Cleveland.Util (runGen)

import Test.Util.Parser

(?==) :: HasCallStack => Text -> Word -> Assertion
codeText ?== expectedCost = do
  parsed <- Parser.codeEntry `shouldParse` ("{" <> codeText <> "}")
  let code = expandList parsed
  expandedInstrsOpSize code @?= OpSize expectedCost

{- All particular numbers below are got using dump-op-size.py script -}

test_Nesting :: [TestTree]
test_Nesting =
  [ testCase "empty string" $ "" ?== 0
  , testCase "{}" $ "{}" ?== 5
  , testCase "{}; {}" $ "{}; {}" ?== 10
  , testCase "100 * {}" $ mconcat (replicate 100 "{}; ") ?== 500
  , testCase "{{}}" $ "{{}}" ?== 10
  ]

test_Numbers :: [TestTree]
test_Numbers =
  [ numTestCase 0 8
  , numTestCase 1 8
  , numTestCase 2 8
  , numTestCase 63 8
  , numTestCase 64 9
  , numTestCase 65 9
  , numTestCase 2048 9
  , numTestCase 20480 10
  , numTestCase (-1) 8
  , numTestCase (-63) 8
  , numTestCase (-64) 9
  ]
  where
  numTestCase (n :: Int) gas =
    testCase (show n) $ ("PUSH int " <> show n <> "; DROP") ?== gas

test_Values :: [TestTree]
test_Values =
  [ testGroup "int"
    [ valueTestCaseExt "small" "int" "5" 8
    , valueTestCaseExt "bigger" "int" "63" 8
    , valueTestCaseExt "big" "int" "64" 9
    ]
  , valueTestCase "nat" "5" 8
  , testGroup "string"
    [ stringTestCase 0 11
    , stringTestCase 1 12
    , stringTestCase 2 13
    , stringTestCase 7 18
    , stringTestCase 8 19
    ]
  , testGroup "bytes"
    [ bytesTestCase 0 11
    , bytesTestCase 1 12
    , bytesTestCase 2 13
    , bytesTestCase 7 18
    , bytesTestCase 8 19
    ]
  , valueTestCase "mutez" "5" 8
  , valueTestCase "bool" "True" 8
  , valueTestCase "key_hash" (dquotes $ pretty (gen genKeyHash)) 47
  , testGroup "timestamp"
    [ valueTestCaseExt "small" "timestamp" "5" 8
    , valueTestCaseExt "bigger" "timestamp" "123" 9
    , valueTestCaseExt "textual" "timestamp" (dquotes "2018-08-08 00:00:00Z") 31
    ]
  , testGroup "address"
    [ valueTestCaseExt "plain tz address" "address" (dquotes $ pretty (gen genAddress)) 47
    , valueTestCaseExt "plain KT address" "address" (dquotes $ pretty (gen genAddress)) 47
    , let addr = gen genAddress
      in valueTestCaseExt "empty entrypoint" "address" (dquotes $ pretty addr <> "%") 48
    , let epAddr = EpAddress (gen genAddress) . unsafe $ buildEpName "a"
      in valueTestCaseExt "short entrypoint" "address" (dquotes $ pretty epAddr) 49
    , let epAddr = EpAddress (gen genAddress) . unsafe . buildEpName . mconcat $ replicate 8 "a"
      in valueTestCaseExt "long entrypoint" "address" (dquotes $ pretty epAddr) 56
    ]
  -- ed25519 and secp256k1 keys have different size
  , valueTestCase "key" (dquotes $ pretty (PublicKeyEd25519 (gen Ed25519.genPublicKey))) 65
  , valueTestCase "key" (dquotes $ pretty (PublicKeySecp256k1 (gen Secp256k1.genPublicKey))) 66
  , valueTestCase "unit" "Unit" 8
  , valueTestCase "signature" (dquotes $ pretty (SignatureEd25519 (gen Ed25519.genSignature))) 110
  , testGroup "chain_id"
    [ valueTestCase "chain_id" (dquotes $ pretty dummyChainId) 26
    , valueTestCase "chain_id" "0x00000000" 15
    ]
  , testGroup "option"
    [ valueTestCaseExt "none" "(option int)" "None" 10
    , valueTestCaseExt "some" "(option int)" "(Some 5)" 12
    ]
  , testGroup "list"
    [ valueTestCaseExt "Empty" "(list int)" "{}" 13
    , valueTestCaseExt "Length 1" "(list int)" "{1}" 15
    , valueTestCaseExt "Length 8" "(list int)" "{1;1;1;1;1;1;1;1}" 29
    ]
  , testGroup "set"
    [ valueTestCaseExt "Empty" "(set int)" "{}" 13
    , valueTestCaseExt "Length 1" "(set int)" "{1}" 15
    , valueTestCaseExt "Length 1" "(set int)" "{1;2;3;4;5;6;7;8}" 29
    ]
  , valueTestCaseExt "pair" "(pair int int)" "(Pair 1 2)" 16
  , valueTestCaseExt "or" "(or int int)" "(Left 1)" 14
  , testGroup "lambda"
    [ valueTestCaseExt "empty lambda" "(lambda int int)" "{}" 15
    , valueTestCaseExt "simple lambda" "(lambda int int)" "{DUP; DROP}" 19
    ]
  , testGroup "map"
    [ valueTestCaseExt "Empty" "(map int int)" "{}" 15
    , valueTestCaseExt "Length 1" "(map int int)" "{Elt 1 2}" 21
    , valueTestCaseExt "Length 1 (big key)" "(map int int)" "{Elt 100 2}" 22
    , let val = "{Elt 1 1; Elt 2 2; Elt 3 3; Elt 4 4; \
                \ Elt 5 5; Elt 6 6; Elt 7 7; Elt 8 8 }"
      in valueTestCaseExt "Length 8" "(map int int)" val 63
    ]
  ]
  where
    valueTestCaseExt name ty val gas =
      testCase name $
        ("PUSH " <> ty <> " " <> val <> "; DROP") ?== gas
    valueTestCase ty val gas =
      valueTestCaseExt (toString ty) ty val gas
    stringTestCase l gas =
      let name = "Length " <> show l
          val = dquotes $ T.replicate l "a"
      in valueTestCaseExt name "string" val gas
    bytesTestCase l gas =
      let name = "Length " <> show l
          val = "0x" <> T.replicate l "12"
      in valueTestCaseExt name "bytes" val gas

    gen :: Gen a -> a
    gen genA = runGen 20 12 genA

test_Types :: [TestTree]
test_Types =
  [ typeTestCase "int" 8
  , typeTestCase "string" 8
  , typeTestCase "mutez" 8
  , typeTestCase "key_hash" 8
  , typeTestCase "address" 8
  , typeTestCase "signature" 8
  , typeTestCase "option int" 10
  , typeTestCase "list int" 10
  , typeTestCase "set int" 10
  , typeTestCase "contract int" 10
  , typeTestCase "pair int int" 12
  , typeTestCase "or int int" 12
  , typeTestCase "or (pair nat nat) (pair int nat)" 20
  , typeTestCase "lambda int unit" 12
  , typeTestCase "map int int" 12
  , typeTestCase "big_map int int" 12
  , typeTestCase "lambda operation int" 12

  , typeTestCase "pair int int" 12
  , typeTestCase "pair (int : %) int" 12
  , typeTestCase "pair (int :a) int" 18
  , typeTestCase "pair (int %a) int" 18
  , typeTestCase "pair (int %a :a) int" 21
  , typeTestCase "pair (int %a) (int %a)" 24
  , typeTestCase "pair :a (int %a) (int %a)" 30
  , typeTestCase "pair :a int (int %a)" 24
  , typeTestCase "pair : (int %a :) (int %a)" 24
  , typeTestCase "pair :a (int %a :a) (int %a)" 33
  ]
  where
  typeTestCase t gas =
    testCase (toString t) $ ("DUP; CONTRACT " <> t <> "; DROP") ?== gas

test_Instructions :: [TestTree]
test_Instructions =
  [ instrTestCase "FAILWITH" 2
  , instrTestCase "FAIL" 9

  , instrTestCaseFailWith "DUP" 4
  , instrTestCaseFailWith "DUP @" 4
  , instrTestCaseFailWith "DUP @a" 10
  , instrTestCaseFailWith "DUP; DROP" 6
  , instrTestCaseFailWith "DUP; SWAP" 6
  , instrTestCaseFailWith "DUP; SWAP; SWAP" 8
  , instrTestCaseFailWith "PUSH int 0" 8
  , instrTestCaseFailWith "PUSH @a int 0" 14
  , instrTestCaseFailWith "PUSH @a (int :a) 0" 20
  , instrTestCaseFailWith "SOME" 4
  , instrTestCaseFailWith "SOME @" 4
  , instrTestCaseFailWith "SOME @a" 10
  , instrTestCaseFailWith "NONE int" 6
  , instrTestCaseFailWith "NONE (option int)" 8
  , instrTestCaseFailWith "NONE :a @a (option :a int)" 23
  , instrTestCaseFailWith "UNIT" 4
  , instrTestCaseFailWith "DUP; SOME; IF_NONE {}{DROP}" 20
  , instrTestCaseFailWith "DUP; DUP; PAIR" 8
  , instrTestCaseFailWith "DUP; DUP; PAIR % % : @" 8
  , instrTestCaseFailWith "DUP; DUP; PAIR %a" 14
  , instrTestCaseFailWith "DUP; DUP; PAIR %a %a" 17
  , instrTestCaseFailWith "DUP; DUP; PAIR %a %a :a" 20
  , instrTestCaseFailWith "DUP; DUP; PAIR %a % :a @a" 20
  , instrTestCaseFailWith "DUP; DUP; PAIR % %a :a @a" 23
  , instrTestCaseFailWith "DUP; DUP; PAIR %a %a :a @a" 23
    -- Further skipping some instructions since everything is trivial
  , instrTestCaseFailWith "EMPTY_BIG_MAP :a @a (int :a) (unit :a)" 29
  , instrTestCaseFailWith "PUSH (list int) {}; MAP {}" 20
  , instrTestCaseFailWith "DIP {}" 9
  , instrTestCaseFailWith "DIP 0 {}" 11
  , instrTestCaseFailWith "DIP 1 {}" 11
  , instrTestCaseFailWith "DROP 0" 6
  , instrTestCaseFailWith "PACK" 4
  , instrTestCaseFailWith "PUSH int 1; PUSH int 2; ADD" 16
  , instrTestCaseFailWith "DUP; CONTRACT int" 8

  , namedInstrTestCaseFailWith "CREATE_CONTRACT"
      "UNIT; PUSH mutez 5; NONE key_hash; \
      \CREATE_CONTRACT \
        \{ parameter unit; \
        \  storage unit; \
        \  code { CDR; NIL operation; PAIR } }" 44

  , namedInstrTestCaseFailWith "CREATE_CONTRACT + PUSH or"
      "PUSH (or int nat) (Left 5); PUSH mutez 5; NONE key_hash; \
      \CREATE_CONTRACT \
        \{ parameter unit; \
        \  storage (or int nat); \
        \  code { DROP; PUSH (or int nat) (Right 2); NIL operation; PAIR } }" 70
  , namedInstrTestCaseFailWith "CREATE_CONTRACT + PUSH or + Annotation"
      "PUSH (or %root int nat) (Left 5); PUSH mutez 5; NONE key_hash; \
      \CREATE_CONTRACT \
        \{ parameter unit; \
        \  storage (or int nat); \
        \  code { DROP; PUSH (or int nat) (Right 2); NIL operation; PAIR } }" 70
  ]
  where
  namedInstrTestCase testName instr gas =
    testCase testName $ instr ?== gas

  namedInstrTestCaseFailWith testName instr gas =
    namedInstrTestCase testName (instr <> "; FAILWITH") gas

  instrTestCase instr gas =
    namedInstrTestCase (toString instr) instr gas

  instrTestCaseFailWith instr gas =
    instrTestCase (instr <> "; FAILWITH") gas