-- | Functions to import contracts to be used in tests.

module Michelson.Test.Import
  (
    -- * Read, parse, typecheck
    readContract
  , importContract
  , importUntypedContract
  , ImportContractError (..)

    -- * Tasty helpers
  , testTreesWithContract
  , testTreesWithTypedContract
  , testTreesWithUntypedContract
  , concatTestTrees

    -- * HSpec helpers
  , specWithContract
  , specWithTypedContract
  , specWithUntypedContract
  ) where

import Control.Exception (IOException)
import Data.Singletons (SingI, demote)
import Data.Typeable ((:~:)(..), eqT)
import Fmt (Buildable(build), pretty, (+|), (|+))
import Test.Hspec (Spec, describe, expectationFailure, it, runIO)
import Test.HUnit (assertFailure)
import Test.Tasty (TestTree)
import Test.Tasty.HUnit (testCase)

import Michelson.Parser.Error (ParserException(..))
import Michelson.Runtime (parseExpandContract, prepareContract)
import Michelson.TypeCheck (SomeContract(..), TCError, typeCheckContract)
import Michelson.Typed (Contract, FullContract(..), toUType)
import qualified Michelson.Untyped as U
import Util.IO

----------------------------------------------------------------------------
-- tasty helpers
----------------------------------------------------------------------------

-- | Import contract and use to create test trees. Both versions of contract are
-- passed to the callback function (untyped and typed).
--
-- If contract's import fails, a tree with single failing test will be generated
-- (so test tree will likely be generated unexceptionally, but a failing
-- result will notify about problem).
testTreesWithContract
  :: (Each [Typeable, SingI] [cp, st], HasCallStack)
  => FilePath -> ((U.Contract, Contract cp st) -> IO [TestTree]) -> IO [TestTree]
testTreesWithContract = testTreesWithContractImpl importContract

-- | Like 'testTreesWithContract' but supplies only untyped contract.
testTreesWithUntypedContract
  :: HasCallStack
  => FilePath -> (U.Contract -> IO [TestTree]) -> IO [TestTree]
testTreesWithUntypedContract =
  testTreesWithContractImpl importUntypedContract

-- | Like 'testTreesWithContract' but supplies only typed contract.
testTreesWithTypedContract
  :: (Each [Typeable, SingI] [cp, st], HasCallStack)
  => FilePath -> (Contract cp st -> IO [TestTree]) -> IO [TestTree]
testTreesWithTypedContract =
  testTreesWithContractImpl (fmap snd . importContract)

testTreesWithContractImpl
  :: HasCallStack
  => (FilePath -> IO contract)
  -> FilePath
  -> (contract -> IO [TestTree])
  -> IO [TestTree]
testTreesWithContractImpl doImport file testImpl =
  saferImport doImport file >>= \case
    Left err -> pure [testCase ("Import contract " <> file) $ assertFailure err]
    Right contract -> testImpl contract

-- A helper function which allows you to use multiple
-- 'testTreesWithTypedContract' in a single top-level test with type
-- 'IO [TestTree]'.
concatTestTrees :: [IO [TestTree]] -> IO [TestTree]
concatTestTrees = fmap concat . sequence

----------------------------------------------------------------------------
-- hspec helpers
----------------------------------------------------------------------------

-- | Import contract and use it in the spec. Both versions of contract are
-- passed to the callback function (untyped and typed).
--
-- If contract's import fails, a spec with single failing expectation
-- will be generated (so tests will likely run unexceptionally, but a failing
-- result will notify about problem).
specWithContract
  :: (Each [Typeable, SingI] [cp, st], HasCallStack)
  => FilePath -> ((U.Contract, Contract cp st) -> Spec) -> Spec
specWithContract = specWithContractImpl importContract

-- | A version of 'specWithContract' which passes only the typed
-- representation of the contract.
specWithTypedContract
  :: (Each [Typeable, SingI] [cp, st], HasCallStack)
  => FilePath -> (Contract cp st -> Spec) -> Spec
specWithTypedContract = specWithContractImpl (fmap snd . importContract)

specWithUntypedContract :: FilePath -> (U.Contract -> Spec) -> Spec
specWithUntypedContract = specWithContractImpl importUntypedContract

specWithContractImpl
  :: HasCallStack
  => (FilePath -> IO contract) -> FilePath -> (contract -> Spec) -> Spec
specWithContractImpl doImport file execSpec =
  either errorSpec (describe ("Test contract " <> file) . execSpec)
    =<< runIO (saferImport doImport file)
  where
    errorSpec = it ("Import contract " <> file) . expectationFailure

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

-- Catch some errors during contract import, we don't want the whole
-- test suite to crash if something like that happens.
saferImport :: (FilePath -> IO contract) -> FilePath -> IO (Either String contract)
saferImport doImport file =
  ((Right <$> doImport file)
  `catch` \(e :: ImportContractError) -> pure $ Left $ displayException e)
  `catch` \(e :: IOException) -> pure $ Left $ displayException e

----------------------------------------------------------------------------
-- Reading, parsing, typechecking
----------------------------------------------------------------------------

readContract
  :: forall cp st .
    Each [Typeable, SingI] [cp, st]
  => FilePath
  -> Text
  -> Either ImportContractError (U.Contract, Contract cp st)
readContract filePath txt = do
  contract <- first ICEParse $ parseExpandContract (Just filePath) txt
  SomeContract (FullContract (instr :: Contract cp' st') _ _)
    <- first ICETypeCheck $ typeCheckContract mempty contract
  case (eqT @cp @cp', eqT @st @st') of
    (Just Refl, Just Refl) -> pure (contract, instr)
    (Nothing, _) -> Left $
      ICEUnexpectedParamType (U.para contract) (toUType $ demote @cp)
    _ -> Left (ICEUnexpectedStorageType (U.stor contract) (toUType $ demote @st))

-- | Import contract from a given file path.
--
-- This function reads file, parses and type checks a contract.
-- Within the typechecking we assume that no contracts are originated,
-- otherwise a type checking error will be caused.
--
-- This function may throw 'IOException' and 'ImportContractError'.
importContract
  :: forall cp st .
     Each [Typeable, SingI] [cp, st]
  => FilePath -> IO (U.Contract, Contract cp st)
importContract file = either throwM pure =<< readContract file <$> readFileUtf8 file

importUntypedContract :: FilePath -> IO U.Contract
importUntypedContract = prepareContract . Just

-- | Error type for 'importContract' function.
data ImportContractError
  = ICEUnexpectedParamType U.Type U.Type
  | ICEUnexpectedStorageType U.Type U.Type
  | ICEParse ParserException
  | ICETypeCheck TCError
  deriving stock (Show, Eq)

instance Buildable ImportContractError where
  build =
    \case
      ICEUnexpectedParamType actual expected ->
        "Unexpected parameter type: " +| actual |+
        ", expected: " +| expected |+ ""
      ICEUnexpectedStorageType actual expected ->
        "Unexpected storage type: " +| actual |+
        ", expected: " +| expected |+ ""
      ICEParse e -> "Failed to parse the contract: " +| e |+ ""
      ICETypeCheck e -> "The contract is ill-typed: " +| e |+ ""

instance Exception ImportContractError where
  displayException = pretty