import Test.Tasty (defaultMain, testGroup, TestTree)
import Test.Tasty.HUnit (Assertion, assertEqual, assertFailure, testCase)

import Codec.Encryption.OpenPGP.ASCIIArmor (decode, decodeLazy, decodeLazyWith, encode, encodeLazy, encodeLazyWith, multipartMerge)
import Codec.Encryption.OpenPGP.ASCIIArmor.Types
import Codec.Encryption.OpenPGP.ASCIIArmor.Utils

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Char8 as BC8
import qualified Data.ByteString.Lazy.Char8 as BLC8
import Data.Digest.CRC24 (crc24)
import Data.List (find)
import Data.Word (Word32)

legacyDecodeOptions :: DecodeOptions
legacyDecodeOptions = defaultDecodeOptions { decodeConformanceProfile = Strict_4880_AA }

legacyEncodeOptions :: EncodeOptions
legacyEncodeOptions = defaultEncodeOptions { encodeConformanceProfile = Strict_4880_AA }

modernDecodeOptions :: DecodeOptions
modernDecodeOptions = defaultDecodeOptions { decodeConformanceProfile = Strict_9580_AA }

testCRC24 :: ByteString -> Word32 -> Assertion
testCRC24 bs crc = assertEqual "crc24" crc (crc24 bs)

testArmorDecode :: FilePath -> [FilePath] -> Assertion
testArmorDecode fp targets = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    tbss <- mapM (\target -> BL.readFile $ "tests/data/" ++ target) targets
    case decodeLazy bs of
        Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
        Right as -> assertEqual ("for " ++ fp) tbss (map getPayload as)
    where
        getPayload (Armor _ _ pl) = pl
        getPayload _ = error "This should not happen."

testArmorMultipartDecode :: FilePath -> FilePath -> Assertion
testArmorMultipartDecode fp target = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    tbs <- BL.readFile $ "tests/data/" ++ target
    case decodeLazy bs of
        Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
        Right as -> assertEqual ("for " ++ fp) tbs (getPayload (multipartMerge as))
    where
        getPayload (Armor _ _ pl) = pl
        getPayload _ = error "This should not happen."

testClearsignedDecodeBody :: FilePath -> FilePath -> Assertion
testClearsignedDecodeBody fp target = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    tbs <- BL.readFile $ "tests/data/" ++ target
    case decodeLazy bs of
        Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
        Right [a] -> assertEqual ("for " ++ fp) (convertEndings tbs) (getBody a)
        _ -> assertFailure "This shouldn't happen."
    where
        getBody (ClearSigned _ txt _) = txt
        getBody _ = error "This should not happen."
        convertEndings = crlfUnlinesLazy . BLC8.lines

testClearsignedDecodeSig :: FilePath -> FilePath -> Assertion
testClearsignedDecodeSig fp target = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    tbs <- BL.readFile $ "tests/data/" ++ target
    case decodeLazy bs of
        Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
        Right [a] -> assertEqual ("for " ++ fp) tbs (getSig a)
        _ -> assertFailure "This shouldn't happen."
    where
        getSig (ClearSigned _ _ (Armor _ _ sig)) = sig
        getSig _ = error "This should not happen."

testLegacyArmorEncode :: [FilePath] -> FilePath -> Assertion
testLegacyArmorEncode fps target = do
    bss <- mapM (\fp -> BL.readFile $ "tests/data/" ++ fp) fps
    tbs <- BL.readFile $ "tests/data/" ++ target
    assertEqual "literaldata" tbs (encodeLazyWith legacyEncodeOptions (map (Armor ArmorMessage [("Version","OpenPrivacy 0.99")]) bss))

testLegacyClearsignedEncode :: FilePath -> FilePath -> FilePath -> Assertion
testLegacyClearsignedEncode ftxt fsig ftarget = do
    txt <- BL.readFile $ "tests/data/" ++ ftxt
    sig <- BL.readFile $ "tests/data/" ++ fsig
    target <- BL.readFile $ "tests/data/" ++ ftarget
    assertEqual "clearsigned encode" target (encodeLazyWith legacyEncodeOptions [ClearSigned [("Hash","SHA1")] txt (Armor ArmorSignature [("Version","OpenPrivacy 0.99")] sig)])

testStrictDecode :: FilePath -> Assertion
testStrictDecode fp = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    assertEqual "strict decode" (decodeLazy bs :: Either String [Armor]) (decode (B.concat . BL.toChunks $ bs) :: Either String [Armor])

testStrictEncode :: FilePath -> Assertion
testStrictEncode fp = do
    bs <- BL.readFile $ "tests/data/" ++ fp
    let fakearmors = [Armor ArmorMessage [("Version","OpenPrivacy 0.99")] bs, ClearSigned [("Hash","SHA1")] bs (Armor ArmorSignature [("Version","OpenPrivacy 0.99")] bs)]
    assertEqual "strict encode" (encodeLazy fakearmors) (BL.fromChunks [encode fakearmors])

testModernDecodeAllowsMissingCRC :: Assertion
testModernDecodeAllowsMissingCRC = do
    bs <- BL.readFile "tests/data/msg1.asc"
    let noCrc = removeLinePrefix (BLC8.pack "=") bs
    case decodeLazyWith defaultDecodeOptions noCrc of
        Left e -> assertFailure $ "modern decode should allow missing CRC24: " ++ e
        Right _ -> return ()

testLegacyDecodeRejectsMissingCRC :: Assertion
testLegacyDecodeRejectsMissingCRC = do
    bs <- BL.readFile "tests/data/msg1.asc"
    let noCrc = removeLinePrefix (BLC8.pack "=") bs
    case (decodeLazyWith legacyDecodeOptions noCrc :: Either String [Armor]) of
        Left _ -> return ()
        Right [] -> return ()
        Right _ -> assertFailure "legacy decode should reject missing CRC24"

testModernDecodeAllowsMalformedCRC :: Assertion
testModernDecodeAllowsMalformedCRC = do
    bs <- BL.readFile "tests/data/msg1.asc"
    let malformed = replaceLinePrefix (BLC8.pack "=") (BLC8.pack "=??!!") bs
    case decodeLazyWith defaultDecodeOptions malformed of
        Left e -> assertFailure $ "modern decode should allow malformed CRC24: " ++ e
        Right _ -> return ()

testModernDecodeAllowsMismatchedCRC :: Assertion
testModernDecodeAllowsMismatchedCRC = do
    bs <- BL.readFile "tests/data/msg1.asc"
    let mismatch = replaceLinePrefix (BLC8.pack "=") (BLC8.pack "=AAAA") bs
    case decodeLazyWith defaultDecodeOptions mismatch of
        Left e -> assertFailure $ "modern decode should allow mismatched CRC24: " ++ e
        Right _ -> return ()

testDecodeIgnoresBase64Whitespace :: Assertion
testDecodeIgnoresBase64Whitespace = do
    bs <- BL.readFile "tests/data/msg1.asc"
    let whitespacey = addWhitespaceToBase64 bs
    assertEqual "whitespace-tolerant decode" (decodeLazy bs :: Either String [Armor]) (decodeLazy whitespacey :: Either String [Armor])

testModernCleartextHeaderStrictness :: Assertion
testModernCleartextHeaderStrictness = do
    bs <- BL.readFile "tests/data/msg3.asc"
    let withVersionHeader = insertHeaderAfterPgpSigned (BLC8.pack "Version: Test 1.0") bs
    case decodeLazyWith legacyDecodeOptions withVersionHeader :: Either String [Armor] of
        Left e -> assertFailure $ "legacy decode should allow non-Hash cleartext headers: " ++ e
        Right as
            | any isClearSigned as -> return ()
            | otherwise -> assertFailure "legacy decode should parse cleartext message"
    case decodeLazyWith modernDecodeOptions withVersionHeader :: Either String [Armor] of
        Left _ -> return ()
        Right as
            | any isClearSigned as -> assertFailure "modern decode should reject non-Hash cleartext headers"
            | otherwise -> return ()
  where
    isClearSigned (ClearSigned _ _ _) = True
    isClearSigned _ = False

testEncodeModernOmitsCRC24 :: Assertion
testEncodeModernOmitsCRC24 = do
    bs <- BL.readFile "tests/data/msg1.gpg"
    let encoded = encodeLazyWith defaultEncodeOptions [Armor ArmorMessage [("Version","OpenPrivacy 0.99")] bs]
    case find (BL.isPrefixOf (BLC8.pack "=")) (BLC8.lines encoded) of
        Nothing -> return ()
        Just _ -> assertFailure "modern encode should omit CRC24 line"

testEncodeLegacyIncludesCRC24 :: Assertion
testEncodeLegacyIncludesCRC24 = do
    bs <- BL.readFile "tests/data/msg1.gpg"
    let encoded = encodeLazyWith legacyEncodeOptions [Armor ArmorMessage [("Version","OpenPrivacy 0.99")] bs]
    case find (BL.isPrefixOf (BLC8.pack "=")) (BLC8.lines encoded) of
        Nothing -> assertFailure "legacy encode should include CRC24 line"
        Just _ -> return ()

removeLinePrefix :: BL.ByteString -> BL.ByteString -> BL.ByteString
removeLinePrefix prefix = BLC8.unlines . filter (not . BL.isPrefixOf prefix) . BLC8.lines

replaceLinePrefix :: BL.ByteString -> BL.ByteString -> BL.ByteString -> BL.ByteString
replaceLinePrefix prefix replacement = BLC8.unlines . map replace . BLC8.lines
    where
        replace line
            | BL.isPrefixOf prefix line = replacement
            | otherwise = line

insertHeaderAfterPgpSigned :: BL.ByteString -> BL.ByteString -> BL.ByteString
insertHeaderAfterPgpSigned header bs =
    case BLC8.lines bs of
        [] -> bs
        (x:xs) -> BLC8.unlines (x : header : xs)

addWhitespaceToBase64 :: BL.ByteString -> BL.ByteString
addWhitespaceToBase64 = BLC8.unlines . map tweak . BLC8.lines
    where
        tweak line
            | BL.null line = line
            | BL.isPrefixOf (BLC8.pack "-----") line = line
            | BLC8.elem ':' line = line
            | BL.isPrefixOf (BLC8.pack "=") line = BL.append (BLC8.pack "= ") (BL.drop 1 line)
            | otherwise = BL.append (BLC8.pack " ") line

tests :: TestTree
tests = testGroup "openpgp-asciiarmor" [
   testGroup "CRC24" [
      testCase "CRC24: A" (testCRC24 (BC8.pack "A") 16680698)
    , testCase "CRC24: Haskell" (testCRC24 (BC8.pack "Haskell") 15612750)
    , testCase "CRC24: hOpenPGP and friends" (testCRC24 (BC8.pack "hOpenPGP and friends") 11940960)
    ]
 , testGroup "ASCII armor" [
      testCase "Decode sample armor" (testArmorDecode "msg1.asc" ["msg1.gpg"])
    , testCase "Decode sample armor with cruft" (testArmorDecode "msg1a.asc" ["msg1.gpg"])
    , testCase "Decode multiple sample armors" (testArmorDecode "msg1b.asc" ["msg1.gpg","msg1.gpg","msg1.gpg"])
    , testCase "Decode detached signature" (testArmorDecode "msg4.asc" ["msg4.sig"])
    , testCase "Decode multi-part armor" (testArmorMultipartDecode "msg2.asc" "msg2.pgp")
    , testCase "Decode body of clear-signed" (testClearsignedDecodeBody "msg3.asc" "msg3")
    , testCase "Decode sig of clear-signed" (testClearsignedDecodeSig "msg3.asc" "msg3.sig")
    , testCase "Encode (legacy) sample armor" (testLegacyArmorEncode ["msg1.gpg"] "msg1.asc")
    , testCase "Encode (legacy) multiple sample armors" (testLegacyArmorEncode ["msg1.gpg","msg1.gpg","msg1.gpg"] "msg1c.asc")
    , testCase "Encode (legacy) clear-signed sig" (testLegacyClearsignedEncode "msg3" "msg3.sig" "msg3.asc")
    , testCase "Decode from strict ByteString" (testStrictDecode "msg1.asc")
    , testCase "Encode to strict ByteString" (testStrictEncode "msg1.gpg")
    , testCase "Decode sample RFC9580 armor" (testArmorDecode "seipdv2.pgp.aa" ["seipdv2.pgp"])
    , testCase "Modern decode accepts missing CRC24" testModernDecodeAllowsMissingCRC
    , testCase "Legacy decode rejects missing CRC24" testLegacyDecodeRejectsMissingCRC
    , testCase "Modern decode accepts malformed CRC24" testModernDecodeAllowsMalformedCRC
    , testCase "Modern decode accepts mismatched CRC24" testModernDecodeAllowsMismatchedCRC
    , testCase "Decode ignores base64 whitespace" testDecodeIgnoresBase64Whitespace
    , testCase "Modern cleartext profile rejects non-Hash headers" testModernCleartextHeaderStrictness
    , testCase "Modern encode omits CRC24" testEncodeModernOmitsCRC24
    , testCase "Legacy encode includes CRC24" testEncodeLegacyIncludesCRC24
    ]
 ]

main :: IO ()
main = defaultMain tests
