import Test.Tasty.QuickCheck
import Test.Tasty

import Control.Applicative
import Data.ASN1.Get (runGet, Result(..))
import Data.ASN1.BitArray
import Data.ASN1.Prim
import Data.ASN1.Serialize
import Data.ASN1.BinaryEncoding.Parse
import Data.ASN1.BinaryEncoding.Writer
import Data.ASN1.BinaryEncoding
import Data.ASN1.Encoding
import Data.ASN1.Types
import Data.ASN1.Types.Lowlevel

import Data.Hourglass

import qualified Data.ByteString as B

import Control.Monad

instance Arbitrary ASN1Class where
        arbitrary = elements [ Universal, Application, Context, Private ]

instance Arbitrary ASN1Length where
        arbitrary = do
                c <- choose (0,2) :: Gen Int
                case c of
                        0 -> liftM LenShort (choose (0,0x79))
                        1 -> do
                                nb <- choose (0x80,0x1000)
                                return $ mkSmallestLength nb
                        _ -> return LenIndefinite
                where
                        nbBytes nb = if nb > 255 then 1 + nbBytes (nb `div` 256) else 1

arbitraryDefiniteLength :: Gen ASN1Length
arbitraryDefiniteLength = arbitrary `suchThat` (\l -> l /= LenIndefinite)

arbitraryTag :: Gen ASN1Tag
arbitraryTag = choose(1,10000)

instance Arbitrary ASN1Header where
        arbitrary = liftM4 ASN1Header arbitrary arbitraryTag arbitrary arbitrary

arbitraryEvents :: Gen ASN1Events
arbitraryEvents = do
        hdr@(ASN1Header _ _ _ len) <- liftM4 ASN1Header arbitrary arbitraryTag (return False) arbitraryDefiniteLength
        let blen = case len of
                LenLong _ x -> x
                LenShort x  -> x
                _           -> 0
        pr <- liftM Primitive (arbitraryBSsized blen)
        return (ASN1Events [Header hdr, pr])

newtype ASN1Events = ASN1Events [ASN1Event]

instance Show ASN1Events where
        show (ASN1Events x) = show x

instance Arbitrary ASN1Events where
        arbitrary = arbitraryEvents


arbitraryOID :: Gen OID
arbitraryOID = do
        i1  <- choose (0,2) :: Gen Integer
        i2  <- choose (0,39) :: Gen Integer
        ran <- choose (0,30) :: Gen Int
        l   <- replicateM ran (suchThat arbitrary (\i -> i > 0))
        return $ (i1:i2:l)

arbitraryBSsized :: Int -> Gen B.ByteString
arbitraryBSsized len = do
        ws <- replicateM len (choose (0, 255) :: Gen Int)
        return $ B.pack $ map fromIntegral ws

instance Arbitrary B.ByteString where
        arbitrary = do
                len <- choose (0, 529) :: Gen Int
                arbitraryBSsized len

instance Arbitrary BitArray where
        arbitrary = do
                bs <- arbitrary
                w  <- choose (0,7) :: Gen Int
                return $ toBitArray bs w

instance Arbitrary Date where
    arbitrary = do
        y <- choose (1951, 2050)
        m <- elements [ January .. December]
        d <- choose (1, 30)
        return $ normalizeDate $ Date y m d

normalizeDate :: Date -> Date
normalizeDate origDate
    | y < 1951  = normalizeDate (Date (y + 50) m d)
    | otherwise = normalizedDate
  where
    normalizedDate@(Date y m d) = timeConvert (timeConvert origDate :: Elapsed)

instance Arbitrary TimeOfDay where
    arbitrary = do
        h    <- choose (0, 23)
        mi   <- choose (0, 59)
        se   <- choose (0, 59)
        nsec <- return 0
        return $ TimeOfDay (Hours h) (Minutes mi) (Seconds se) nsec

instance Arbitrary DateTime where
    arbitrary = DateTime <$> arbitrary <*> arbitrary

instance Arbitrary TimezoneOffset where
    arbitrary = elements [ timezone_UTC, TimezoneOffset 60, TimezoneOffset 120, TimezoneOffset (-360) ]

instance Arbitrary Elapsed where
    arbitrary = Elapsed . Seconds <$> arbitrary

instance Arbitrary ASN1TimeType where
    arbitrary = elements [TimeUTC, TimeGeneralized]

instance Arbitrary ASN1StringEncoding where
    arbitrary = elements [UTF8, Numeric, Printable, T61, VideoTex, IA5, Graphic, Visible, General, UTF32, BMP]

arbitraryPrintString encoding = do
    let printableString = (['a'..'z'] ++ ['A'..'Z'] ++ ['0'..'9'] ++ " ()+,-./:=?")
    asn1CharacterString encoding <$> replicateM 21 (elements printableString)

arbitraryBS encoding = ASN1CharacterString encoding . B.pack <$> replicateM 7 (choose (0,0xff))

arbitraryIA5String = asn1CharacterString IA5 <$> replicateM 21 (choose (toEnum 0,toEnum 127))

arbitraryUCS2 :: Gen ASN1CharacterString
arbitraryUCS2 = asn1CharacterString BMP <$> replicateM 12 (choose (toEnum 0,toEnum 0xffff))

arbitraryUnicode :: ASN1StringEncoding -> Gen ASN1CharacterString
arbitraryUnicode e = asn1CharacterString e <$> replicateM 35 (choose (toEnum 0,toEnum 0x10ffff))

instance Arbitrary ASN1CharacterString where
    arbitrary = oneof
            [ arbitraryUnicode UTF8
            , arbitraryUnicode UTF32
            , arbitraryUCS2
            , arbitraryPrintString Numeric
            , arbitraryPrintString Printable
            , arbitraryBS T61
            , arbitraryBS VideoTex
            , arbitraryIA5String
            , arbitraryPrintString Graphic
            , arbitraryPrintString Visible
            , arbitraryPrintString General
            ]

instance Arbitrary ASN1 where
        arbitrary = oneof
                [ liftM Boolean arbitrary
                , liftM IntVal arbitrary
                , liftM BitString arbitrary
                , liftM OctetString arbitrary
                , return Null
                , liftM OID arbitraryOID
                --, Real Double
                -- , return Enumerated
                , ASN1String <$> arbitrary
                , ASN1Time <$> arbitrary <*> arbitrary <*> arbitrary
                ]

newtype ASN1s = ASN1s [ASN1]

instance Show ASN1s where
        show (ASN1s x) = show x

instance Arbitrary ASN1s where
        arbitrary = do
                x <- choose (0,5) :: Gen Int
                z <- case x of
                        4 -> makeList Sequence
                        3 -> makeList Set
                        _ -> resize 2 $ listOf1 arbitrary
                return $ ASN1s z
                where
                        makeList str = do
                                (ASN1s l) <- arbitrary
                                return ([Start str] ++ l ++ [End str])

prop_header_marshalling_id :: ASN1Header -> Bool
prop_header_marshalling_id v = (ofDone $ runGet getHeader $ putHeader v) == Right v
    where ofDone (Done r _ _) = Right r
          ofDone _            = Left "not done"

prop_event_marshalling_id :: ASN1Events -> Bool
prop_event_marshalling_id (ASN1Events e) = (parseLBS $ toLazyByteString e) == Right e

prop_asn1_der_marshalling_id v = (decodeASN1 DER . encodeASN1 DER) v `assertEq` Right v
    where assertEq got expected
                 | got /= expected = error ("got: " ++ show got ++ " expected: " ++ show expected)
                 | otherwise       = True

marshallingTests = testGroup "Marshalling"
    [ testProperty "Header" prop_header_marshalling_id
    , testProperty "Event"  prop_event_marshalling_id
    , testProperty "DER"    prop_asn1_der_marshalling_id
    ]

main = defaultMain $ testGroup "asn1-encoding" [marshallingTests]