{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE RecordWildCards #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Cursor.Tree.Gen
    (
    ) where

import qualified Data.List.NonEmpty as NE
import Data.Maybe

import Data.GenValidity
import Data.GenValidity.Containers ()

import Test.QuickCheck

import Cursor.Tree

instance GenUnchecked TreeCursorSelection

instance GenValid TreeCursorSelection where
    genValid = genValidStructurally
    shrinkValid = shrinkValidStructurally

instance GenUnchecked a => GenUnchecked (SwapResult a)

instance GenValid a => GenValid (SwapResult a) where
    genValid = genValidStructurally
    shrinkValid = shrinkValidStructurally

instance GenUnchecked a => GenUnchecked (PromoteElemResult a)

instance GenValid a => GenValid (PromoteElemResult a) where
    genValid = genValidStructurally
    shrinkValid = shrinkValidStructurally

instance GenUnchecked a => GenUnchecked (PromoteResult a)

instance GenValid a => GenValid (PromoteResult a) where
    genValid = genValidStructurally
    shrinkValid = shrinkValidStructurally

instance GenUnchecked a => GenUnchecked (DemoteResult a)

instance GenValid a => GenValid (DemoteResult a) where
    genValid = genValidStructurally
    shrinkValid = shrinkValidStructurally

instance GenUnchecked a => GenUnchecked (CTree a) where
    genUnchecked =
        sized $ \n -> do
            s <- upTo n
            (a, b) <- genSplit s
            val <- resize a genUnchecked
            for <- resize b genUnchecked
            pure $ CNode val for
    shrinkUnchecked (CNode a cf) =
        [CNode a' cf' | (a', cf') <- shrinkUnchecked (a, cf)]

instance GenValid a => GenValid (CTree a) where
    genValid =
        sized $ \n -> do
            s <- upTo n
            (a, b) <- genSplit s
            val <- resize a genValid
            for <- resize b genValid
            pure $ CNode val for
    shrinkValid (CNode a cf) = [CNode a' cf' | (a', cf') <- shrinkValid (a, cf)]

instance GenUnchecked a => GenUnchecked (CForest a) where
    genUnchecked =
        sized $ \n ->
            case n of
                0 -> pure EmptyCForest
                _ ->
                    oneof
                        [ ClosedForest <$> resize n genUnchecked
                        , OpenForest <$> resize n genUnchecked
                        ]
    shrinkUnchecked EmptyCForest = []
    shrinkUnchecked (ClosedForest ne) =
        EmptyCForest : (ClosedForest <$> shrinkUnchecked ne)
    shrinkUnchecked (OpenForest ne) =
        EmptyCForest :
        ClosedForest (NE.map rebuildCTree ne) :
        (OpenForest <$> shrinkUnchecked ne)

instance GenValid a => GenValid (CForest a) where
    genValid =
        sized $ \n ->
            case n of
                0 -> pure EmptyCForest
                _ ->
                    oneof
                        [ ClosedForest <$> resize n genValid
                        , OpenForest <$> resize n genValid
                        ]
    shrinkValid EmptyCForest = []
    shrinkValid (ClosedForest ne) =
        EmptyCForest : (ClosedForest <$> shrinkValid ne)
    shrinkValid (OpenForest ne) =
        EmptyCForest :
        ClosedForest (NE.map rebuildCTree ne) : (OpenForest <$> shrinkValid ne)

instance (GenUnchecked a, GenUnchecked b) => GenUnchecked (TreeCursor a b) where
    genUnchecked =
        sized $ \n -> do
            s <- upTo n
            (a, b, c, d) <- genSplit4 s
            treeAbove <- resize (a + b) genUnchecked
            treeCurrent <- resize c genUnchecked
            treeBelow <- resize d genUnchecked
            pure TreeCursor {..}
    shrinkUnchecked tc =
        let opts =
                catMaybes
                    [ do ta <- treeAbove tc
                         pure $ tc {treeAbove = treeAboveAbove ta}
                    ]
         in opts ++ genericShrinkUnchecked tc

instance (GenValid a, GenValid b) => GenValid (TreeCursor a b) where
    genValid =
        sized $ \n -> do
            s <- upTo n
            (a, b, c, d) <- genSplit4 s
            treeAbove <- resize (a + b) genValid
            treeCurrent <- resize c genValid
            treeBelow <- resize d genValid
            pure TreeCursor {..}
    shrinkValid = shrinkValidStructurallyWithoutExtraFiltering

instance GenUnchecked b => GenUnchecked (TreeAbove b) where
    genUnchecked =
        sized $ \n -> do
            s <- upTo n
            (a, b, c, d) <- genSplit4 s
            treeAboveLefts <- resize a genUnchecked
            treeAboveAbove <- resize b genUnchecked
            treeAboveNode <- resize c genUnchecked
            treeAboveRights <- resize d genUnchecked
            pure TreeAbove {..}

instance GenValid b => GenValid (TreeAbove b) where
    genValid =
        sized $ \n -> do
            s <- upTo n
            (a, b, c, d) <- genSplit4 s
            treeAboveLefts <- resize a genValid
            treeAboveAbove <- resize b genValid
            treeAboveNode <- resize c genValid
            treeAboveRights <- resize d genValid
            pure TreeAbove {..}
    shrinkValid = shrinkValidStructurallyWithoutExtraFiltering