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

import Cartel
import Cartel.Betsy
import Cartel.Ast
import Cartel.Render
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
import Data.List (nub, sort)

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
                , fmap ltEq genVersion
                , fmap gtEq 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 (\(Flag nm _) -> nm) 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 defaultExtensions neslist
  , fmap otherExtensions 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
  allNames <- replicateM (fromIntegral nFlags) genNonEmptyString
  let names = nub . sort $ allNames
      mkf nm = do
        opts <- arbitrary
        _ <- lift (makeFlag nm opts)
        return ()
  _ <- mapM mkf names
  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
    Cabal <$> arbitrary <*> arbitrary <*> arbitrary <*>
      lift currentFlags

-- | 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'
--
-- * runs the 'Betsy', which results in a State computation with the
-- 'StdGen' as the state
--
-- * runs the state computation with an initial state initialized by
-- the given seed
--
-- * fetches just the result of the state computation, discarding the
-- ending state
--
-- * fetches just the resulting Cabal record, discarding the flags
-- created in the 'Betsy'
--
-- * Renders either the error or the Cabal record
genCabalText
  :: Word
  -- ^ Size
  -> Int
  -- ^ Random seed
  -> Either String String
  -- ^ Either an error message, or the Cabal text.
genCabalText sz sd
  = either (Left . renderNoIndent) (Right . renderNoIndent)
  . fmap fst
  . fst
  . flip runState (mkStdGen sd)
  . runBetsy
  . flip runReaderT sz
  $ (arbitrary :: Gen Cabal)


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