content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE NoFieldSelectors #-}
  
  module Skeletest.Internal.Snapshot (
    -- * Running snapshot
    SnapshotContext (..),
    SnapshotResult (..),
    updateSnapshot,
    checkSnapshot,
  
    -- * Rendering
    SnapshotRenderer (..),
    defaultSnapshotRenderers,
    setSnapshotRenderers,
    getSnapshotRenderers,
    plainRenderer,
    renderWithShow,
  
    -- ** SnapshotFile
    SnapshotFile (..),
    SnapshotValue (..),
    decodeSnapshotFile,
    encodeSnapshotFile,
    normalizeSnapshotFile,
  
    -- * Infrastructure
    getAndIncSnapshotIndex,
    SnapshotUpdateFlag (..),
  ) where
  
  import Control.Monad.IO.Class (MonadIO)
  import Control.Monad.Trans.Except (runExceptT, throwE)
  import Data.Aeson qualified as Aeson
  import Data.Aeson.Encode.Pretty qualified as Aeson
  import Data.Char (isAlpha, isPrint)
  import Data.Map.Strict (Map)
  import Data.Map.Strict qualified as Map
  import Data.Maybe (fromMaybe, mapMaybe)
  import Data.Text (Text)
  import Data.Text qualified as Text
  import Data.Text.IO qualified as Text
  import Data.Text.Lazy qualified as TextL
  import Data.Text.Lazy.Encoding qualified as TextL
  import Data.Typeable (Typeable)
  import Data.Typeable qualified as Typeable
  import Data.Void (absurd)
  import Debug.RecoverRTTI (anythingToString)
  import System.Directory (createDirectoryIfMissing)
  import System.FilePath (replaceExtension, splitFileName, takeDirectory, (</>))
  import System.IO.Error (isDoesNotExistError)
  import System.IO.Unsafe (unsafePerformIO)
  import UnliftIO.Exception (throwIO, try)
  import UnliftIO.IORef (
    IORef,
    atomicModifyIORef',
    modifyIORef',
    newIORef,
    readIORef,
    writeIORef,
   )
  
  import Skeletest.Internal.CLI (FlagSpec (..), IsFlag (..))
  import Skeletest.Internal.Error (SkeletestError (..), invariantViolation)
  import Skeletest.Internal.Fixtures (
    Fixture (..),
    FixtureScope (..),
    getFixture,
    noCleanup,
    withCleanup,
   )
  import Skeletest.Internal.TestInfo (TestInfo (..), getTestInfo)
  import Skeletest.Internal.Utils.Map qualified as Map.Utils
  
  {----- Infrastructure -----}
  
  data SnapshotTestFixture = SnapshotTestFixture
    { snapshotIndexRef :: IORef Int
    }
  
  instance Fixture SnapshotTestFixture where
    fixtureAction = do
      snapshotIndexRef <- newIORef 0
      pure . noCleanup $ SnapshotTestFixture{..}
  
  getAndIncSnapshotIndex :: (MonadIO m) => m Int
  getAndIncSnapshotIndex = do
    SnapshotTestFixture{snapshotIndexRef} <- getFixture
    atomicModifyIORef' snapshotIndexRef $ \i -> (i + 1, i)
  
  data SnapshotFileFixture = SnapshotFileFixture
    { snapshotFileRef :: IORef (Maybe SnapshotFile)
    }
  
  instance Fixture SnapshotFileFixture where
    fixtureScope = PerFileFixture
    fixtureAction = do
      TestInfo{testFile} <- getTestInfo
      let snapshotPath = getSnapshotPath testFile
  
      mSnapshotFile <-
        try (Text.readFile snapshotPath) >>= \case
          Left e
            | isDoesNotExistError e -> pure Nothing
            | otherwise -> throwIO e
          Right contents ->
            case decodeSnapshotFile contents of
              Just snapshotFile -> pure $ Just snapshotFile
              Nothing -> throwIO $ SnapshotFileCorrupted snapshotPath
      let snapshotChanged newSnapshot = mSnapshotFile /= Just newSnapshot
  
      snapshotFileRef <- newIORef mSnapshotFile
      pure . withCleanup SnapshotFileFixture{..} $
        -- write snapshot back out when file is done
        readIORef snapshotFileRef >>= \case
          Just snapshotFile | snapshotChanged snapshotFile -> do
            createDirectoryIfMissing True (takeDirectory snapshotPath)
            Text.writeFile snapshotPath $ encodeSnapshotFile $ normalizeSnapshotFile snapshotFile
          _ -> pure ()
  
  newtype SnapshotUpdateFlag = SnapshotUpdateFlag Bool
  
  instance IsFlag SnapshotUpdateFlag where
    flagName = "update"
    flagShort = Just 'u'
    flagHelp = "Update snapshots"
    flagSpec = SwitchFlag SnapshotUpdateFlag
  
  {----- Running snapshot -----}
  
  data SnapshotContext = SnapshotContext
    { snapshotRenderers :: [SnapshotRenderer]
    , snapshotTestInfo :: TestInfo
    , snapshotIndex :: Int
    }
  
  updateSnapshot :: (Typeable a, MonadIO m) => SnapshotContext -> a -> m ()
  updateSnapshot snapshotContext testResult = do
    SnapshotFileFixture{snapshotFileRef} <- getFixture
    modifyIORef' snapshotFileRef (Just . setSnapshot . fromMaybe emptySnapshotFile)
    where
      SnapshotContext
        { snapshotRenderers = renderers
        , snapshotTestInfo = testInfo@TestInfo{testModule}
        , snapshotIndex
        } = snapshotContext
  
      emptySnapshotFile =
        SnapshotFile
          { moduleName = testModule
          , snapshots = Map.empty
          }
  
      testIdentifier = toTestIdentifier testInfo
      renderedTestResult = renderVal renderers testResult
      setSnapshot snapshotFile@SnapshotFile{snapshots} =
        let setForTest = Map.Utils.adjustNested (setAt snapshotIndex renderedTestResult) testIdentifier
         in snapshotFile{snapshots = setForTest snapshots}
  
      -- Set the given snapshot at the given index. If the index is too large,
      -- fill in with empty snapshots.
      --
      -- >>> setAt 3 "x" ["a"] == ["a", "", "", "x"]
      setAt i0 v =
        let go = \cases
              i [] -> replicate i emptySnapshotVal <> [v]
              0 (_ : xs) -> v : xs
              i (x : xs) -> x : go (i - 1) xs
         in if i0 < 0
              then invariantViolation $ "Got negative snapshot index: " <> show i0
              else go i0
      emptySnapshotVal = SnapshotValue{snapshotContent = "", snapshotLang = Nothing}
  
  data SnapshotResult
    = SnapshotMissing
    | SnapshotMatches
    | SnapshotDiff
        { snapshotContent :: Text
        , renderedTestResult :: Text
        }
    deriving (Show, Eq)
  
  checkSnapshot :: (Typeable a, MonadIO m) => SnapshotContext -> a -> m SnapshotResult
  checkSnapshot snapshotContext testResult =
    fmap (either id absurd) . runExceptT $ do
      SnapshotFileFixture{snapshotFileRef} <- getFixture
      fileSnapshots <-
        readIORef snapshotFileRef >>= \case
          Nothing -> returnE SnapshotMissing
          Just SnapshotFile{snapshots} -> pure snapshots
  
      let snapshots = Map.Utils.findOrEmpty (toTestIdentifier testInfo) fileSnapshots
      snapshot <- maybe (returnE SnapshotMissing) pure $ safeIndex snapshots snapshotIndex
  
      let (snapshotContent, renderedTestResult) = (getContent snapshot, getContent renderedTestResultVal)
      returnE $
        if snapshotContent == renderedTestResult
          then SnapshotMatches
          else SnapshotDiff{snapshotContent, renderedTestResult}
    where
      SnapshotContext
        { snapshotRenderers = renderers
        , snapshotTestInfo = testInfo
        , snapshotIndex
        } = snapshotContext
  
      returnE = throwE
      renderedTestResultVal = renderVal renderers testResult
  
      safeIndex xs0 i0 =
        let go = \cases
              _ [] -> Nothing
              0 (x : _) -> Just x
              i (_ : xs) -> go (i - 1) xs
         in if i0 < 0 then Nothing else go i0 xs0
  
  {----- Snapshot file -----}
  
  data SnapshotFile = SnapshotFile
    { moduleName :: Text
    , snapshots :: Map TestIdentifier [SnapshotValue]
    -- ^ full test identifier => snapshots
    -- e.g. ["group1", "group2", "returns val1 and val2"] => ["val1", "val2"]
    }
    deriving (Show, Eq)
  
  data SnapshotValue = SnapshotValue
    { snapshotContent :: Text
    , snapshotLang :: Maybe Text
    }
    deriving (Show, Eq)
  
  getContent :: SnapshotValue -> Text
  getContent SnapshotValue{snapshotContent} = snapshotContent
  
  type TestIdentifier = [Text]
  
  getSnapshotPath :: FilePath -> FilePath
  getSnapshotPath testFile = testDir </> "__snapshots__" </> snapshotFileName
    where
      (testDir, testFileName) = splitFileName testFile
      snapshotFileName = replaceExtension testFileName ".snap.md"
  
  toTestIdentifier :: TestInfo -> TestIdentifier
  toTestIdentifier TestInfo{testContexts, testName} = testContexts <> [testName]
  
  decodeSnapshotFile :: Text -> Maybe SnapshotFile
  decodeSnapshotFile = parseFile . Text.lines
    where
      parseFile = \case
        line : rest
          | Just moduleName <- Text.stripPrefix "# " line -> do
              let snapshotFile =
                    SnapshotFile
                      { moduleName = Text.strip moduleName
                      , snapshots = Map.empty
                      }
              parseSections snapshotFile Nothing rest
        _ -> Nothing
  
      parseSections ::
        SnapshotFile
        -- \^ The parsed snapshot file so far
        -> Maybe [Text]
        -- \^ The current test identifier, if one is set
        -> [Text]
        -- \^ The rest of the lines to process
        -> Maybe SnapshotFile
      parseSections snapshotFile@SnapshotFile{snapshots} mTest = \case
        [] -> pure snapshotFile
        line : rest
          -- ignore empty lines
          | "" <- Text.strip line -> parseSections snapshotFile mTest rest
          -- found a test section
          | Just sectionName <- Text.stripPrefix "## " line -> do
              let testIdentifier = map Text.strip $ Text.splitOn " / " sectionName
              let snapshotFile' = snapshotFile{snapshots = Map.insert testIdentifier [] snapshots}
              parseSections snapshotFile' (Just testIdentifier) rest
          -- found the beginning of a snapshot
          | Just lang <- (Text.stripPrefix "```" . Text.strip) line -> do
              testIdentifier <- mTest
              (snapshot, rest') <- parseSnapshot [] rest
              let
                snapshotVal =
                  SnapshotValue
                    { snapshotContent = snapshot
                    , snapshotLang = if Text.null lang then Nothing else Just lang
                    }
                snapshotFile' = snapshotFile{snapshots = Map.adjust (<> [snapshotVal]) testIdentifier snapshots}
              parseSections snapshotFile' mTest rest'
          -- anything else is invalid
          | otherwise -> Nothing
  
      parseSnapshot snapshot = \case
        [] -> Nothing
        line : rest
          | "```" <- Text.strip line -> pure (Text.unlines snapshot, rest)
          | otherwise -> parseSnapshot (snapshot <> [line]) rest
  
  encodeSnapshotFile :: SnapshotFile -> Text
  encodeSnapshotFile SnapshotFile{..} =
    Text.intercalate "\n" $
      h1 moduleName : concatMap toSection (Map.toList snapshots)
    where
      toSection (testIdentifier, snaps) =
        h2 (Text.intercalate " / " testIdentifier) : map codeBlock snaps
  
      h1 s = "# " <> s <> "\n"
      h2 s = "## " <> s <> "\n"
      codeBlock SnapshotValue{..} =
        Text.concat
          [ "```" <> fromMaybe "" snapshotLang <> "\n"
          , snapshotContent
          , "```\n"
          ]
  
  normalizeSnapshotFile :: SnapshotFile -> SnapshotFile
  normalizeSnapshotFile file@SnapshotFile{snapshots} =
    file
      { snapshots = Map.fromList . map normalize . Map.toList $ snapshots
      }
    where
      normalize (testIdentifier, vals) =
        ( map (sanitizeNonPrint . sanitizeSlashes . Text.strip) testIdentifier
        , map normalizeSnapshotVal vals
        )
  
      sanitizeSlashes = Text.replace " /" " \\/"
  
      sanitizeNonPrint = Text.concatMap $ \case
        c | (not . isPrint) c -> Text.drop 1 . Text.dropEnd 1 . Text.pack . show $ c
        c -> Text.singleton c
  
  {----- Renderers -----}
  
  data SnapshotRenderer
    = forall a.
    (Typeable a) =>
    SnapshotRenderer
    { render :: a -> Text
    , snapshotLang :: Maybe Text
    }
  
  plainRenderer :: (Typeable a) => (a -> Text) -> SnapshotRenderer
  plainRenderer render =
    SnapshotRenderer
      { render
      , snapshotLang = Nothing
      }
  
  renderWithShow :: forall a. (Typeable a, Show a) => SnapshotRenderer
  renderWithShow = plainRenderer (Text.pack . show @a)
  
  defaultSnapshotRenderers :: [SnapshotRenderer]
  defaultSnapshotRenderers =
    [ plainRenderer @String Text.pack
    , plainRenderer @Text id
    , jsonRenderer
    ]
    where
      jsonRenderer =
        SnapshotRenderer
          { render = TextL.toStrict . TextL.decodeUtf8 . Aeson.encodePretty @Aeson.Value
          , snapshotLang = Just "json"
          }
  
  renderVal :: (Typeable a) => [SnapshotRenderer] -> a -> SnapshotValue
  renderVal renderers a =
    normalizeSnapshotVal $
      case mapMaybe tryRender renderers of
        [] ->
          SnapshotValue
            { snapshotContent = Text.pack $ anythingToString a
            , snapshotLang = Nothing
            }
        rendered : _ -> rendered
    where
      tryRender SnapshotRenderer{..} =
        let toValue v = SnapshotValue{snapshotContent = render v, snapshotLang}
         in toValue <$> Typeable.cast a
  
  normalizeSnapshotVal :: SnapshotValue -> SnapshotValue
  normalizeSnapshotVal SnapshotValue{..} =
    SnapshotValue
      { snapshotContent = normalizeTrailingNewlines snapshotContent
      , snapshotLang = collapse $ Text.filter isAlpha <$> snapshotLang
      }
    where
      collapse = \case
        Just "" -> Nothing
        m -> m
  
      -- Ensure there's exactly one trailing newline.
      normalizeTrailingNewlines s = Text.dropWhileEnd (== '\n') s <> "\n"
  
  snapshotRenderersRef :: IORef [SnapshotRenderer]
  snapshotRenderersRef = unsafePerformIO $ newIORef []
  {-# NOINLINE snapshotRenderersRef #-}
  
  setSnapshotRenderers :: [SnapshotRenderer] -> IO ()
  setSnapshotRenderers = writeIORef snapshotRenderersRef
  
  getSnapshotRenderers :: (MonadIO m) => m [SnapshotRenderer]
  getSnapshotRenderers = readIORef snapshotRenderersRef
  