module Main where

import Codec.Compression.Zstd.FFI qualified as Zstd
import Control.Exception (bracket)
import Control.Monad (guard, unless, void)
import Data.ByteString qualified as BS
import Data.ByteString.Builder (toLazyByteString)
import Data.ByteString.Lazy qualified as BSL
import Data.Foldable (for_)
import Data.List (isSuffixOf)
import Data.Traversable (for)
import Data.Vector qualified as Vector
import Debug.Trace (traceM)
import Foreign qualified
import Shower qualified
import System.Directory (listDirectory)
import System.Environment (lookupEnv)
import System.FilePath ((</>))
import Test.Tasty (testGroup, defaultMain)
import Test.Tasty.HUnit (testCase)

import Codec.Ktx qualified as Ktx1
import Codec.Ktx.KeyValue qualified as KeyValue
import Codec.Ktx2.DFD (DFD(..))
import Codec.Ktx2.DFD.Khronos.BasicV2 qualified as BasicV2
import Codec.Ktx2.Header qualified as Header
import Codec.Ktx2.Level qualified as Level
import Codec.Ktx2.Read qualified as Read
import Codec.Ktx2.Write qualified as Write

assetsPath :: FilePath
assetsPath = ".." </> "assets"

main :: IO ()
main = do
  assets <- listDirectory assetsPath
  defaultMain $
    testGroup "ktx-codec"
      [ testGroup "KTX 1" do
          fp <- assets
          guard $ ".ktx" `isSuffixOf` fp
          pure . testCase fp $
            testKtx1 (assetsPath </> fp)
      , testGroup "KTX 2 (files)" do
          fp <- assets
          guard $ ".ktx2" `isSuffixOf` fp
          pure . testCase fp $
            testKtx2File (assetsPath </> fp)
      , testGroup "KTX 2 (bytes)" do
          fp <- assets
          guard $ ".ktx2" `isSuffixOf` fp
          pure . testCase fp $
            testKtx2Bytes (assetsPath </> fp)
      ]

testKtx1 :: FilePath -> IO ()
testKtx1 source =
  Ktx1.fromFile source >>= \case
    Left (offset, err) -> do
      traceM err
      fail $ unwords
        [ "KTX-1 load error in"
        , source
        , "at"
        , show offset
        ]
    Right ktx -> do
      printer ktx
      let encodedBytes = toLazyByteString (Ktx1.toBuilder ktx)
      originalBytes <- BSL.readFile source
      unless (encodedBytes == originalBytes) $
        fail $ "Encoded KTX does not match original for " <> source

testKtx2File :: FilePath -> IO ()
testKtx2File source =
  bracket (Read.open source) Read.close $
    void . testKtx2Context source

testKtx2Bytes :: FilePath -> IO ()
testKtx2Bytes source = do
  contents <- BS.readFile source
  context <- Read.bytes contents

  chunks <- testKtx2Context source context
  -- BSL.writeFile (source <> ".out") chunks

  tripping <- Read.bytes (BSL.toStrict chunks)
  let
    readHeader = Read.header context
    writeHeader = Read.header tripping
  unless (readHeader == writeHeader) $
    fail . mappend "Header tripping failed:\n" $ show
      ( readHeader
      , writeHeader
      )

  readLevels <- Read.levels context
  writeLevels <- Read.levels tripping
  unless (readLevels == writeLevels) $
    fail . mappend "Level index tripping failed:\n" $ show
      ( readLevels
      , writeLevels
      )

  readDfd <- Read.dataFormatDescriptor context
  writeDfd <- Read.dataFormatDescriptor tripping
  unless (readDfd == writeDfd) $
    fail . mappend "DFD tripping failed:\n" $ show
      ( readDfd
      , writeDfd
      )

  readKvd <- Read.keyValueData context
  writeKvd <- Read.keyValueData tripping
  unless (readKvd == writeKvd) $
    fail . mappend "KeyValue tripping failed:\n" $ show
      ( readKvd
      , writeKvd
      )

  for_ readLevels \level -> do
    readLevel <- Read.levelData context level
    writeLevel <- Read.levelData tripping level
    unless (readLevel == writeLevel) $
      fail . mappend "Level tripping failed:\n" $ show
        ( level
        , BS.take 16 readLevel
        , BS.take 16 writeLevel
        )
testKtx2Context
  :: ( Read.ReadChunk a
     , Read.ReadLevel a
     , Show (Read.Context a)
     )
  => FilePath
  -> Read.Context a
  -> IO BSL.ByteString
testKtx2Context source ctx = do
  printer ctx

  levels <- Read.levels ctx
  printer levels

  DFD{dfdBlocks} <- Read.dataFormatDescriptor ctx
  Vector.forM_ dfdBlocks \block -> do
    case BasicV2.fromBlock block of
      Nothing -> do
        printer block
        fail $ "Unexpected block type in test suite in " <> source
      Just basic -> do
        -- printer basic
        let back = BasicV2.toBlock basic
        unless (back == block) do
          printer (back, block)
          fail "BasicV2 DFD tripping failed"

  kvd <- Read.keyValueData ctx
  printer $ KeyValue.textual kvd

  sgd <- Read.supercompressionGlobalData ctx
  printer sgd

  levelBuffers <- for (Vector.toList levels) \level -> do
    bytes <- Read.levelData ctx level
    unless (BS.length bytes == fromIntegral (Level.byteLength level)) $
      fail "Uncompressed level data does not match with index!"

    Foreign.allocaBytesAligned (fromIntegral $ Level.uncompressedByteLength level) 16 \dst -> do
      case Header.supercompressionScheme (Read.header ctx) of
        Header.SC_NONE -> do
          -- XXX: Write to destination directly
          ok <- Read.levelToPtr ctx level dst
          unless ok $
            fail $ "Data read failed for " <> show level

        Header.SC_ZSTANDARD ->
          -- XXX: Prepare an intermediate buffer for decompression
          Foreign.allocaBytesAligned (fromIntegral $ Level.byteLength level) 16 \stage -> do
            ok <- Read.levelToPtr ctx level stage
            unless ok $
              fail $ "Data read failed for " <> show level

            let expected = fromIntegral $ Level.uncompressedByteLength level
            res <-
              Zstd.checkError $
                Zstd.decompress
                  dst
                  expected
                  stage
                  (fromIntegral $ Level.byteLength level)
            case res of
              Right decoded | decoded == expected ->
                printer (decoded, level)
              Right mismatch ->
                fail $ "Zstd decoded unexpected number of bytes (" <> show mismatch <> ") for level " <> show level
              Left err ->
                fail $ "Zstd error: " <> err

        huh ->
          fail $ "Unexpected supercompression: " <> show huh
    pure
      ( if Header.supercompressionScheme (Read.header ctx) == Header.SC_NONE then
          Nothing
        else
          Just . fromIntegral $ Level.uncompressedByteLength level
      , bytes
      )

  pure $
    Write.toChunks
      (Read.header ctx)
      dfdBlocks
      kvd
      sgd
      levelBuffers

printer :: Show a => a -> IO ()
printer x =
  lookupEnv "TEST_DUMP" >>= \case
    Nothing ->
      pure ()
    Just _ ->
      Shower.printer x