{-# LANGUAGE TypeSynonymInstances #-}
module Cartel.Generators where

import Cartel
import Cartel.Betsy
import Control.Applicative
import Control.Monad.Trans.State
import Control.Monad.Trans.Class
import System.Random (StdGen, random, randomR, Random(..), mkStdGen)
import Control.Monad.Trans.Reader
import Control.Monad

type Gen = ReaderT Word (Betsy (State StdGen))

class Arbitrary a where
  arbitrary :: Gen a

getSize :: Gen Word
getSize = ask

resize :: Word -> Gen a -> Gen a
resize w = withReaderT (const w)

getStdGen :: Gen StdGen
getStdGen = lift . lift $ get

setStdGen :: StdGen -> Gen ()
setStdGen = lift . lift . put

choose :: Random a => (a, a) -> Gen a
choose pair = do
  g <- getStdGen
  let (r, g') = randomR pair g
  setStdGen g'
  return r

chooseDefault :: Random a => Gen a
chooseDefault = do
  g <- getStdGen
  let (r, g') = random g
  setStdGen g'
  return r

frequency :: [(Int, Gen a)] -> Gen a
frequency [] = error "frequency used with empty list"
frequency xs0 = choose (1, tot) >>= (`pick` xs0)
  where
    tot = sum (map fst xs0)
    pick n ((k, x):xs)
      | n <= k = x
      | otherwise = pick (n - k) xs
    pick _ _ = error "pick used with empty list"

elements :: [a] -> Gen a
elements [] = error "elements used with empty list"
elements xs = (xs !!) `fmap` choose (0, length xs - 1)

oneOf :: [Gen a] -> Gen a
oneOf = join . elements

instance Arbitrary Bool where
  arbitrary = chooseDefault

instance Arbitrary Word where
  arbitrary = choose (0, 255)

instance Arbitrary a => Arbitrary (Maybe a) where
  arbitrary = frequency [(3, fmap Just arbitrary), (1, return Nothing)]

instance Arbitrary a => Arbitrary [a] where
  arbitrary = do
    sz <- getSize
    c <- choose (0, sz)
    replicateM (fromIntegral c) arbitrary

instance Arbitrary Char where
  arbitrary = choose ('a', 'z')

genNonEmptyString :: Gen NonEmptyString
genNonEmptyString = liftM2 (:) arbitrary (resize 5 arbitrary)

instance Arbitrary Vcs where
  arbitrary = oneOf $ fmap cvs genNonEmptyString
    : map return [ darcs, git, svn, mercurial,
                   bazaar, archVcs, monotone]

instance Arbitrary RepoKind where
  arbitrary = elements [repoHead, repoThis]

instance Arbitrary Repository where
  arbitrary = Repository <$> arbitrary <*> arbitrary <*> arbitrary
    <*> genNonEmptyString <*> arbitrary <*> arbitrary

listOf :: Gen a -> Gen [a]
listOf g = do
  sz <- getSize
  len <- choose (0, sz)
  replicateM (fromIntegral len) g

genNonEmptyList :: Gen a -> Gen [a]
genNonEmptyList g = liftM2 (:) g (listOf g)

genVersion :: Gen Version
genVersion = resize 3 (genNonEmptyList (choose (0, 99)))

instance Arbitrary Constraint where
  arbitrary = do
    sz <- getSize
    lvls <- choose (0, sz)
    go lvls
    where
      go lvl | lvl == 0 = oneOf
                [ fmap lt genVersion
                , fmap gt genVersion
                , fmap eq genVersion
                , return anyVersion
                ]
             | otherwise = do
                l <- next
                r <- next
                op <- elements [ (&&&), (|||) ]
                return $ l `op` r
        where
          next = choose (0, lvl - 1) >>= go

instance Arbitrary Compiler where
  arbitrary = elements [ ghc, nhc, yhc, hugs, helium, jhc, lhc ]

instance Arbitrary Condition where
  arbitrary = do
    sz <- getSize
    lvls <- choose (0, sz)
    goCondition lvls

goCondition :: Word -> Gen Condition
goCondition lvl
  | lvl == 0 = lift currentFlags >>= lowest
  | otherwise = upper
  where
    lowest [] = nonFlagLevel
    lowest xs = elements (fmap fst xs) >>= withFlagLevel
    nonFlagLevel = oneOf [ fmap system genNonEmptyString
                         , fmap arch genNonEmptyString
                         , liftM2 impl arbitrary arbitrary
                         , return true
                         , return false
                         ]
    withFlagLevel fl = frequency
      [ (2, nonFlagLevel), (1, return $ flag fl)]
    upper = frequency [(2, binary), (1, inverted)]
      where
        next = choose (0, lvl - 1) >>= goCondition
        inverted = fmap invert next
        binary = do
          op <- elements [(&&&), (|||)]
          liftM2 op next next

genCondBlock :: (Arbitrary a, HasBuildInfo a) => Gen a
genCondBlock = liftM3 condBlock arbitrary
  (liftM2 (,) arbitrary arbitrary) arbitrary

instance Arbitrary Package where
  arbitrary = liftM2 package genNonEmptyString arbitrary

genBuildInfo :: HasBuildInfo a => Gen a
genBuildInfo = oneOf
  [ return haskell98
  , return haskell2010
  , fmap buildDepends arbitrary
  , fmap otherModules neslist
  , fmap hsSourceDirs neslist
  , fmap extensions neslist
  , fmap buildTools (genNonEmptyList arbitrary)
  , fmap buildable arbitrary
  , fmap ghcOptions neslist
  , fmap ghcProfOptions neslist
  , fmap ghcSharedOptions neslist
  , fmap hugsOptions neslist
  , fmap nhc98Options neslist
  , fmap includes neslist
  , fmap installIncludes neslist
  , fmap includeDirs neslist
  , fmap cSources neslist
  , fmap extraLibraries neslist
  , fmap ccOptions neslist
  , fmap cppOptions neslist
  , fmap ldOptions neslist
  , fmap pkgConfigDepends (genNonEmptyList arbitrary)
  , fmap frameworks neslist
  ]
  where
    neslist = genNonEmptyList genNonEmptyString

genMainIs :: BuildsExe a => Gen a
genMainIs = fmap mainIs genNonEmptyString

instance Arbitrary FlagOpts where
  arbitrary = liftM3 FlagOpts arbitrary arbitrary arbitrary

addFlagsToState
  :: Word
  -- ^ How many flags?
  -> Gen ()
addFlagsToState nFlags = do
  let mkf = do
        nm <- genNonEmptyString
        opts <- arbitrary
        _ <- lift (makeFlag nm opts)
        return ()
  _ <- replicateM (fromIntegral nFlags) mkf
  return ()

downsize :: Gen a -> Gen a
downsize g = do
  sz <- getSize
  resize (sz - 1) g

instance Arbitrary LibraryField where
  arbitrary = getSize >>= f
    where
      nonNestGens =
        [ (6, genBuildInfo)
        , (1, fmap exposed arbitrary)
        , (1, fmap exposedModules (genNonEmptyList genNonEmptyString))
        ]
      f sz | sz == 0 = nonNest
           | otherwise = nest
        where
          nonNest = frequency nonNestGens
          nest = downsize . frequency $
            (1, genCondBlock) : nonNestGens


instance Arbitrary ExecutableField where
  arbitrary = getSize >>= f
    where
      nonNestGens =
        [ (6, genBuildInfo)
        , (1, genMainIs)
        ]
      nestGens = (1, genCondBlock) : nonNestGens
      f sz | sz == 0 = frequency nonNestGens
           | otherwise = downsize $ frequency nestGens

instance Arbitrary TestSuiteField where
  arbitrary = getSize >>= f
    where
      nonNestGens =
        [ (7, genBuildInfo)
        , (1, fmap testModule genNonEmptyString)
        , (1, pure exitcodeStdio)
        , (1, pure detailed)
        , (1, genMainIs)
        ]
      nestGens = (1, genCondBlock) : nonNestGens
      f sz | sz == 0 = frequency nonNestGens
           | otherwise = downsize $ frequency nestGens

instance Arbitrary BenchmarkField where
  arbitrary = getSize >>= f
    where
      nonNest =
        [ (6, genBuildInfo)
        , (1, return exitcodeStdio)
        ]
      nest = (1, genCondBlock) : nonNest
      f sz | sz == 0 = frequency nonNest
           | otherwise = downsize $ frequency nest

instance Arbitrary BuildType where
  arbitrary = elements [simple, configure, make, custom]

instance Arbitrary License where
  arbitrary = elements [ gpl, agpl, lgpl, bsd2, bsd3, bsd4, mit, mpl,
                         apache, publicDomain, allRightsReserved,
                         otherLicense ]

instance (Arbitrary a, Arbitrary b) => Arbitrary (a, b) where
  arbitrary = liftM2 (,) arbitrary arbitrary

instance Arbitrary Properties where
  arbitrary = Properties
    <$> genNonEmptyString -- name
    <*> genVersion -- version
    <*> arbitrary -- cabalVersion
    <*> arbitrary -- buildType
    <*> arbitrary -- license
    <*> arbitrary -- licenseFile
    <*> listOf genNonEmptyString -- licenseFiles
    <*> arbitrary -- copyright
    <*> arbitrary -- author
    <*> arbitrary
    <*> arbitrary
    <*> arbitrary
    <*> arbitrary -- bugReports
    <*> arbitrary
    <*> arbitrary
    <*> arbitrary -- description
    <*> arbitrary
    <*> arbitrary
    <*> listOf genNonEmptyString -- dataFiles
    <*> arbitrary
    <*> listOf genNonEmptyString
    <*> listOf genNonEmptyString
    <*> listOf genNonEmptyString -- extraTmpFiles

instance Arbitrary Section where
  arbitrary = frequency [(3, nonRepo), (1, repo)]
    where
      nonRepo = oneOf
        [ liftM2 executable genNonEmptyString arbitrary
        , liftM2 testSuite genNonEmptyString arbitrary
        , liftM2 benchmark genNonEmptyString arbitrary
        ]
      repo = fmap repository arbitrary

instance Arbitrary Cabal where
  arbitrary = do
    sz <- getSize
    nFlags <- choose (0, sz)
    addFlagsToState nFlags
    liftM3 cabal arbitrary arbitrary arbitrary

-- | Generates Cabal text given a size parameter and int to generate a
-- seed.  Read the function from the bottom up:
--
-- * generates a random computation in ReaderT
--
-- * runs the ReaderT with the given size, resulting in a 'Betsy'
--
-- * renders the 'Betsy', which results in a State computation with
-- the 'StdGen' as the state
--
-- * runs the State computation with an initialized generator,
-- resulting in a pair (a, s), where @s@ is the resulting state and
-- @a@ is an Either
--
-- * gets the Either, discards the resulting state, resulting in an
-- @Either Error String@
--
-- * transforms the @Error@ to a string
genCabalText
  :: Word
  -- ^ Size
  -> Int
  -- ^ Random seed
  -> Either String String
  -- ^ Either an error message, or the Cabal text.
genCabalText sz sd
  = either (Left . renderIndented 0) (Right . id)
  . fst
  . flip runState (mkStdGen sd)
  . renderBetsy
  . flip runReaderT sz
  $ arbitrary

genResult
  :: Gen a
  -- ^ Generator
  -> Word
  -- ^ How many flags?
  -> Word
  -- ^ Size
  -> Int
  -- ^ Random seed
  -> Either String a
genResult gen nFlags sz sd
  = either (Left . renderIndented 0) (Right . fst)
  . fst
  . flip runState (mkStdGen sd)
  . runBetsy
  . flip runReaderT sz
  $ (addFlagsToState nFlags >> gen)