module Main where

import Control.Monad
import qualified Data.Map as Map

import BEncode.BEncode
import Test.QuickCheck
import Test.QuickCheck.Batch

{-
  Test by either loading into ghci and run

  > quickCheck prop_identity

  or run 'main'.

-}

-- Testing using QuickCheck

main :: IO ()
main = do
  runTests "BEncode" defOpt
    [ run prop_identity
    , run prop_get_shallow
    ]

-- | Calculate the depth of a BEncode tree.
depth :: BEncode -> Int
depth (BInt _) = 1
depth (BString _) = 1
depth (BList lst) = 1 + maximum (0:(map depth lst))
depth (BDict mp) = 1 + maximum (0:(map depth (Map.elems mp)))

-- | Count the number of elements in a BEncode tree.
elems :: BEncode -> Int
elems (BInt _) = 1
elems (BString _) = 1
elems (BList lst) = 1 + sum (map elems lst)
elems (BDict mp) = 1 + sum (map elems (Map.elems mp))

instance Arbitrary BEncode where
  arbitrary = join arbitrary

instance Arbitrary (Gen BEncode) where
  arbitrary = sized $ \n -> sizedBE n
    where
    sizedBE n = frequency 
      [ (2, return genBInt)
      , (2, return genBString)
      , (1, return $ genBList n)
      , (1, return $ genBDict n)
      ]

    genBInt :: Gen BEncode
    genBInt = liftM BInt arbitrary

    genBString :: Gen BEncode
    genBString = liftM BString $ elements ["a", "b", "c", "d", "e"]

    -- lists are monotyped
    genBList :: Int -> Gen BEncode
    genBList n | n <= 0 = return (BList [])
	       | otherwise = do
	no_of_elems <- choose (0,n)
	g <- sizedBE (n `div` no_of_elems)
	v <- vectorOf' no_of_elems g
	return (BList v)

    -- dicts are not monotyped
    genBDict :: Int -> Gen BEncode
    genBDict n | n <= 0 = return (BDict Map.empty)
	       | otherwise = do
	no_of_elems <- choose (0,n)
	v <- vectorOf' no_of_elems (join $ sizedBE (n `div` no_of_elems)) 
	k <- vectorOf' no_of_elems $ elements ["a", "b", "c", "d", "e"]
	return (BDict (Map.fromList (zip k v))) 

-- Taken from a unreleased version of quickcheck
-- Just added ' to the names
--   / Kolmodin
listOf' :: Gen a -> Gen [a]
listOf' gen = sized $ \n ->
  do k <- choose (0,n)
     vectorOf' k gen

vectorOf' :: Int -> Gen a -> Gen [a]
vectorOf' k gen = sequence [ gen | _ <- [1..k] ]


-- QUICKCHECK Properties

-- | Throwing a BEncoded string through a show\/read pair should be the
--   identity
prop_identity :: BEncode -> Property
prop_identity be =
  let d = depth be in
  trivial (d == 1) $
  collect d $
  Just be == bRead (bShow be)

instance Arbitrary Char where
  arbitrary = elements ['a', 'b', 'c', 'd', 'e']

prop_get_shallow :: String -> BEncode -> Bool
prop_get_shallow k be = 
  let d = (BDict (Map.fromList [(k,be)])) in
  be == get k d