{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}

import Apecs (Cache, Component (..), Entity (..), Has (..), Map, Not (..), Proxy (..), System, cfold, destroy, exists, get, makeWorld, runSystem, set)
import Apecs.Core (ExplMembers, explMembers)
import Apecs.Stores.Flags (Flags)
import Apecs.Stores.Flags qualified as Flags
import Apecs.Stores.Query (Every (..), Indexes)
import Apecs.Stores.Query qualified as Query
import Apecs.Stores.Roaring (Roaring)
import Control.Monad (forM, forM_, unless)
import Control.Monad.IO.Class (liftIO)
import Data.IntMap.Strict qualified as IntMap
import Data.IntSet qualified as IntSet
import Data.List (sort)
import Data.Vector.Unboxed qualified as Unboxed
import GHC.Generics (Generic)
import System.Exit (exitFailure)
import Test.QuickCheck (Arbitrary (..), Gen, Property, Testable (..), conjoin, counterexample, isSuccess, quickCheckWithResult, stdArgs, (.&&.), (===))
import Test.QuickCheck qualified as QuickCheck
import Test.QuickCheck.Monadic (monadicIO, run)

data Flag = Flag
  deriving (Eq, Show, Generic)

instance Component Flag where
  type Storage Flag = Flags Flag

data Enemy = Enemy
  deriving (Generic)

instance Component Enemy where
  type Storage Enemy = Flags Enemy

data Dead = Dead
  deriving (Generic)

instance Component Dead where
  type Storage Dead = Flags Dead

newtype Health = Health Int
  deriving (Eq, Show)

instance Component Health where
  type Storage Health = Roaring (Cache 2 (Map Health))

makeWorld "World" [''Flag, ''Enemy, ''Dead, ''Health]

entityId :: Gen Int
entityId =
  QuickCheck.frequency
    [ (4, QuickCheck.choose (-16, 64))
    , (1, arbitrary)
    , (1, QuickCheck.elements [minBound, maxBound, -1, -2])
    ]

data Op
  = Raise Int
  | Lower Int
  deriving (Show)

instance Arbitrary Op where
  arbitrary = QuickCheck.oneof [Raise <$> entityId, Lower <$> entityId]

model :: [Op] -> IntSet.IntSet
model = foldl step IntSet.empty
  where
    step flagged = \case
      Raise ety -> IntSet.insert ety flagged
      Lower ety -> IntSet.delete ety flagged

listed :: forall c. (Has World IO c, ExplMembers IO (Storage c)) => System World [Int]
listed = do
  store <- getStore @World @IO @c
  Unboxed.toList <$> liftIO (explMembers store)

prop_flagsMatchModel :: [Op] -> [Int] -> Property
prop_flagsMatchModel ops probes = monadicIO $ do
  (memberIds, probed, folded) <- run $ do
    world <- initWorld
    flip runSystem world $ do
      forM_ ops $ \case
        Raise ety -> set (Entity ety) Flag
        Lower ety -> destroy (Entity ety) $ Proxy @Flag
      memberIds <- listed @Flag
      probed <- forM (IntSet.toList expected <> probes) $ \ety -> (,) ety <$> exists (Entity ety) (Proxy @Flag)
      folded <- cfold (\total Flag -> total + 1) (0 :: Int)
      pure (memberIds, probed, folded)
  pure $
    counterexample "members" (sort memberIds === IntSet.toAscList expected)
      .&&. counterexample "exists" (probed === [(ety, IntSet.member ety expected) | (ety, _) <- probed])
      .&&. counterexample "cfold" (folded === IntSet.size expected)
  where
    expected = model ops

prop_flagsClear :: [Int] -> Property
prop_flagsClear marked = monadicIO $ do
  (counted, memberIds, remaining, restored) <- run $ do
    world <- initWorld
    flip runSystem world $ do
      forM_ marked $ \ety -> set (Entity ety) Flag
      Flags.clear @Flag
      counted <- Query.count @Flag
      memberIds <- listed @Flag
      remaining <- or <$> forM marked (\ety -> exists (Entity ety) $ Proxy @Flag)
      forM_ marked $ \ety -> set (Entity ety) Flag
      restored <- Query.count @Flag
      pure (counted, memberIds, remaining, restored)
  pure $
    counted === 0
      .&&. memberIds === []
      .&&. remaining === False
      .&&. restored === IntSet.size (IntSet.fromList marked)

data Edit
  = SetEnemy Int
  | DropEnemy Int
  | SetDead Int
  | DropDead Int
  | SetHealth Int Int
  | DropHealth Int
  deriving (Show)

instance Arbitrary Edit where
  arbitrary =
    QuickCheck.frequency
      [ (3, SetEnemy <$> cachedId)
      , (1, DropEnemy <$> cachedId)
      , (3, SetDead <$> cachedId)
      , (1, DropDead <$> cachedId)
      , (3, SetHealth <$> cachedId <*> arbitrary)
      , (1, DropHealth <$> cachedId)
      ]
    where
      cachedId = QuickCheck.frequency [(8, QuickCheck.choose (-4, 24)), (1, entityId)] `QuickCheck.suchThat` (/= cacheVacancy)

cacheVacancy :: Int
cacheVacancy = -2

data World' = World'
  { enemies :: IntSet.IntSet
  , dead :: IntSet.IntSet
  , healths :: IntMap.IntMap Int
  }

replay :: [Edit] -> World'
replay = foldl step World'{enemies = IntSet.empty, dead = IntSet.empty, healths = IntMap.empty}
  where
    step known = \case
      SetEnemy ety -> known{enemies = IntSet.insert ety known.enemies}
      DropEnemy ety -> known{enemies = IntSet.delete ety known.enemies}
      SetDead ety -> known{dead = IntSet.insert ety known.dead}
      DropDead ety -> known{dead = IntSet.delete ety known.dead}
      SetHealth ety hp -> known{healths = IntMap.insert ety hp known.healths}
      DropHealth ety -> known{healths = IntMap.delete ety known.healths}

perform :: Edit -> System World ()
perform = \case
  SetEnemy ety -> set (Entity ety) Enemy
  DropEnemy ety -> destroy (Entity ety) $ Proxy @Enemy
  SetDead ety -> set (Entity ety) Dead
  DropDead ety -> destroy (Entity ety) $ Proxy @Dead
  SetHealth ety hp -> set (Entity ety) $ Health hp
  DropHealth ety -> destroy (Entity ety) $ Proxy @Health

data Answer = Answer [Int] Int

answer :: forall c. (Has World IO c, ExplMembers IO (Storage c), Indexes (Storage c)) => System World Answer
answer = Answer <$> listed @c <*> Query.count @c

agrees :: String -> IntSet.IntSet -> Answer -> Property
agrees name expected (Answer memberIds counted) =
  counterexample name $
    counterexample "members" (sort memberIds === IntSet.toAscList expected)
      .&&. counterexample "count" (counted === IntSet.size expected)

prop_roaringMatchesModel :: [Edit] -> [Int] -> Property
prop_roaringMatchesModel edits probes = monadicIO $ do
  (memberIds, values, probed, counted) <- run $ do
    world <- initWorld
    flip runSystem world $ do
      forM_ edits perform
      memberIds <- listed @Health
      values <- forM (IntMap.keys expected) $ \ety -> (,) ety <$> get (Entity ety)
      probed <- forM cachedProbes $ \ety -> (,) ety <$> exists (Entity ety) (Proxy @Health)
      counted <- Query.count @Health
      pure (memberIds, values, probed, counted)
  pure $
    counterexample "members" (sort memberIds === IntMap.keys expected)
      .&&. counterexample "values" (values === [(ety, Health hp) | (ety, hp) <- IntMap.toList expected])
      .&&. counterexample "exists" (probed === [(ety, IntMap.member ety expected) | ety <- cachedProbes])
      .&&. counterexample "count" (counted === IntMap.size expected)
  where
    expected = (replay edits).healths
    cachedProbes = filter (/= cacheVacancy) probes

prop_queriesMatchModel :: [Edit] -> Property
prop_queriesMatchModel edits = monadicIO $ do
  (single, pair, without, triple, mixed, tupled, healthSum) <- run $ do
    world <- initWorld
    flip runSystem world $ do
      forM_ edits perform
      single <- answer @(Every Enemy)
      pair <- answer @(Every (Enemy, Health))
      without <- answer @(Every (Enemy, Not Dead))
      triple <- answer @(Every (Enemy, Health, Dead))
      mixed <- answer @(Every (Enemy, Health, Not Dead))
      tupled <- Query.count @(Enemy, Health)
      healthSum <- cfold (\total (Every (Enemy, Health hp)) -> total + toInteger hp) 0
      pure (single, pair, without, triple, mixed, tupled, healthSum)
  pure
    . QuickCheck.checkCoverage
    . QuickCheck.cover 40 (not $ IntSet.null withHealth) "Enemy and Health overlap"
    . QuickCheck.cover 20 (not . IntSet.null $ withHealth `IntSet.difference` known.dead) "Enemy and Health overlap outside Dead"
    . QuickCheck.cover 10 (not . IntSet.null $ withHealth `IntSet.intersection` known.dead) "Enemy, Health and Dead overlap"
    $ conjoin
      [ agrees "Every Enemy" known.enemies single
      , agrees "Every (Enemy, Health)" withHealth pair
      , agrees "Every (Enemy, Not Dead)" (known.enemies `IntSet.difference` known.dead) without
      , agrees "Every (Enemy, Health, Dead)" (withHealth `IntSet.intersection` known.dead) triple
      , agrees "Every (Enemy, Health, Not Dead)" (withHealth `IntSet.difference` known.dead) mixed
      , counterexample "count (Enemy, Health)" $ tupled === IntSet.size withHealth
      , counterexample "cfold sum" $ healthSum === sum [toInteger hp | (ety, hp) <- IntMap.toList known.healths, IntSet.member ety known.enemies]
      ]
  where
    known = replay edits
    withHealth = known.enemies `IntSet.intersection` IntMap.keysSet known.healths

main :: IO ()
main = do
  results <-
    mapM
      (quickCheckWithResult stdArgs{QuickCheck.maxSuccess = 1000})
      [ property prop_flagsMatchModel
      , property prop_flagsClear
      , property prop_roaringMatchesModel
      , property prop_queriesMatchModel
      ]
  unless (all isSuccess results) exitFailure