-- | Module for custom instance of Data.HashMap.Strict that decodes from any list of objects that have an id field
module Calamity.Internal.SnowflakeMap where

import Calamity.Internal.Utils ()
import Calamity.Types.Snowflake
import Control.DeepSeq
import Control.Lens.At
import Control.Lens.Iso
import Control.Lens.Wrapped
import Data.Aeson (FromJSON (..), ToJSON (..), withArray)
import Data.Data
import qualified Data.Foldable as F
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as SH
import Data.Hashable
import GHC.Exts (IsList)
import GHC.Generics
import TextShow
import qualified TextShow.Generic as TSG
import Unsafe.Coerce

newtype SnowflakeMap a = SnowflakeMap
  { unSnowflakeMap :: HashMap (Snowflake a) a
  }
  deriving (Generic, Eq, Data, Ord, Show)
  deriving (TextShow) via TSG.FromGeneric (SnowflakeMap a)
  deriving newtype (IsList, Semigroup, Monoid)
  deriving newtype (NFData, Hashable)

-- instance At (SnowflakeMap a) where
--   at k f m = at (unSnowflakeMap k) f m

instance Functor SnowflakeMap where
  fmap f = SnowflakeMap . coerceSnowflakeMap . fmap f . unSnowflakeMap

instance Foldable SnowflakeMap where
  foldr f b = Prelude.foldr f b . unSnowflakeMap

instance Traversable SnowflakeMap where
  traverse f = fmap (SnowflakeMap . coerceSnowflakeMap) . traverse f . unSnowflakeMap

-- deriving instance NFData a => NFData (SnowflakeMap a)

-- deriving instance Hashable a => Hashable (SnowflakeMap a)

instance Wrapped (SnowflakeMap a) where
  type Unwrapped (SnowflakeMap a) = HashMap (Snowflake a) a

  _Wrapped' = iso unSnowflakeMap SnowflakeMap

type instance Index (SnowflakeMap a) = Snowflake a

type instance IxValue (SnowflakeMap a) = a

instance SnowflakeMap a ~ t => Rewrapped (SnowflakeMap b) a

instance Ixed (SnowflakeMap a) where
  ix i = _Wrapped . ix i

instance At (SnowflakeMap a) where
  at i = _Wrapped . at i

over :: (HashMap (Snowflake a) a -> HashMap (Snowflake b) b) -> SnowflakeMap a -> SnowflakeMap b
over f = SnowflakeMap . f . unSnowflakeMap
{-# INLINEABLE over #-}

-- SAFETY: 'Snowflake' always uses the underlying hash function (Word64)
coerceSnowflakeMap :: HashMap (Snowflake a) v -> HashMap (Snowflake b) v
coerceSnowflakeMap = unsafeCoerce
{-# INLINEABLE coerceSnowflakeMap #-}

empty :: SnowflakeMap a
empty = SnowflakeMap SH.empty
{-# INLINEABLE empty #-}

singleton :: HasID' a => a -> SnowflakeMap a
singleton v = SnowflakeMap $ SH.singleton (getID v) v
{-# INLINEABLE singleton #-}

null :: SnowflakeMap a -> Bool
null = SH.null . unSnowflakeMap
{-# INLINEABLE null #-}

size :: SnowflakeMap a -> Int
size = SH.size . unSnowflakeMap
{-# INLINEABLE size #-}

member :: Snowflake a -> SnowflakeMap a -> Bool
member k = SH.member k . unSnowflakeMap
{-# INLINEABLE member #-}

lookup :: Snowflake a -> SnowflakeMap a -> Maybe a
lookup k = SH.lookup k . unSnowflakeMap
{-# INLINEABLE lookup #-}

lookupDefault :: a -> Snowflake a -> SnowflakeMap a -> a
lookupDefault d k = SH.lookupDefault d k . unSnowflakeMap
{-# INLINEABLE lookupDefault #-}

(!) :: SnowflakeMap a -> Snowflake a -> a
(!) m k = unSnowflakeMap m SH.! k
{-# INLINEABLE (!) #-}

infixl 9 !

insert :: HasID' a => a -> SnowflakeMap a -> SnowflakeMap a
insert v = over $ SH.insert (getID v) v
{-# INLINEABLE insert #-}

insertWith :: HasID' a => (a -> a -> a) -> a -> SnowflakeMap a -> SnowflakeMap a
insertWith f v = over $ SH.insertWith f (getID v) v
{-# INLINEABLE insertWith #-}

delete :: Snowflake a -> SnowflakeMap a -> SnowflakeMap a
delete k = over $ SH.delete k
{-# INLINEABLE delete #-}

adjust :: (a -> a) -> Snowflake a -> SnowflakeMap a -> SnowflakeMap a
adjust f k = over $ SH.adjust f k
{-# INLINEABLE adjust #-}

update :: (a -> Maybe a) -> Snowflake a -> SnowflakeMap a -> SnowflakeMap a
update f k = over $ SH.update f k
{-# INLINEABLE update #-}

alter :: (Maybe a -> Maybe a) -> Snowflake a -> SnowflakeMap a -> SnowflakeMap a
alter f k = over $ SH.alter f k
{-# INLINEABLE alter #-}

union :: SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
union m m' = SnowflakeMap $ SH.union (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE union #-}

unionWith :: (a -> a -> a) -> SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
unionWith f m m' = SnowflakeMap $ SH.unionWith f (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE unionWith #-}

unionWithKey :: (Snowflake a -> a -> a -> a) -> SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
unionWithKey f m m' = SnowflakeMap $ SH.unionWithKey f (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE unionWithKey #-}

unions :: [SnowflakeMap a] -> SnowflakeMap a
unions = SnowflakeMap . SH.unions . Prelude.map unSnowflakeMap
{-# INLINEABLE unions #-}

map :: (a1 -> a2) -> SnowflakeMap a1 -> SnowflakeMap a2
map f = over $ coerceSnowflakeMap . SH.map f
{-# INLINEABLE map #-}

mapWithKey :: (Snowflake a1 -> a1 -> a2) -> SnowflakeMap a1 -> SnowflakeMap a2
mapWithKey f = over $ coerceSnowflakeMap . SH.mapWithKey f
{-# INLINEABLE mapWithKey #-}

traverseWithKey :: Applicative f => (Snowflake a1 -> a1 -> f a2) -> SnowflakeMap a1 -> f (SnowflakeMap a2)
traverseWithKey f = fmap (SnowflakeMap . coerceSnowflakeMap) . SH.traverseWithKey f . unSnowflakeMap
{-# INLINEABLE traverseWithKey #-}

difference :: SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
difference m m' = SnowflakeMap $ SH.difference (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE difference #-}

differenceWith :: (a -> a -> Maybe a) -> SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
differenceWith f m m' = SnowflakeMap $ SH.differenceWith f (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE differenceWith #-}

intersection :: SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap a
intersection m m' = SnowflakeMap $ SH.intersection (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE intersection #-}

intersectionWith :: (a -> a -> b) -> SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap b
intersectionWith f m m' = SnowflakeMap . coerceSnowflakeMap $ SH.intersectionWith f (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE intersectionWith #-}

intersectionWithKey :: (Snowflake a -> a -> a -> b) -> SnowflakeMap a -> SnowflakeMap a -> SnowflakeMap b
intersectionWithKey f m m' = SnowflakeMap . coerceSnowflakeMap $ SH.intersectionWithKey f (unSnowflakeMap m) (unSnowflakeMap m')
{-# INLINEABLE intersectionWithKey #-}

foldl' :: (a -> b -> a) -> a -> SnowflakeMap b -> a
foldl' f s m = SH.foldl' f s $ unSnowflakeMap m
{-# INLINEABLE foldl' #-}

foldlWithKey' :: (a -> Snowflake b -> b -> a) -> a -> SnowflakeMap b -> a
foldlWithKey' f s m = SH.foldlWithKey' f s $ unSnowflakeMap m
{-# INLINEABLE foldlWithKey' #-}

foldr :: (b -> a -> a) -> a -> SnowflakeMap b -> a
foldr f s m = SH.foldr f s $ unSnowflakeMap m
{-# INLINEABLE foldr #-}

foldrWithKey :: (Snowflake b -> b -> a -> a) -> a -> SnowflakeMap b -> a
foldrWithKey f s m = SH.foldrWithKey f s $ unSnowflakeMap m
{-# INLINEABLE foldrWithKey #-}

filter :: (a -> Bool) -> SnowflakeMap a -> SnowflakeMap a
filter f = over $ SH.filter f
{-# INLINEABLE filter #-}

filterWithKey :: (Snowflake a -> a -> Bool) -> SnowflakeMap a -> SnowflakeMap a
filterWithKey f = over $ SH.filterWithKey f
{-# INLINEABLE filterWithKey #-}

mapMaybe :: (a -> Maybe b) -> SnowflakeMap a -> SnowflakeMap b
mapMaybe f = over $ coerceSnowflakeMap . SH.mapMaybe f
{-# INLINEABLE mapMaybe #-}

mapMaybeWithKey :: (Snowflake a -> a -> Maybe b) -> SnowflakeMap a -> SnowflakeMap b
mapMaybeWithKey f = over $ coerceSnowflakeMap . SH.mapMaybeWithKey f
{-# INLINEABLE mapMaybeWithKey #-}

keys :: SnowflakeMap a -> [Snowflake a]
keys = SH.keys . unSnowflakeMap
{-# INLINEABLE keys #-}

elems :: SnowflakeMap a -> [a]
elems = SH.elems . unSnowflakeMap
{-# INLINEABLE elems #-}

toList :: SnowflakeMap a -> [(Snowflake a, a)]
toList = SH.toList . unSnowflakeMap
{-# INLINEABLE toList #-}

fromList :: HasID' a => [a] -> SnowflakeMap a
fromList = SnowflakeMap . SH.fromList . Prelude.map (\v -> (getID v, v))
{-# INLINEABLE fromList #-}

fromListWith :: HasID' a => (a -> a -> a) -> [a] -> SnowflakeMap a
fromListWith f = SnowflakeMap . SH.fromListWith f . Prelude.map (\v -> (getID v, v))
{-# INLINEABLE fromListWith #-}

instance (FromJSON a, HasID' a) => FromJSON (SnowflakeMap a) where
  parseJSON = withArray "SnowflakeMap" $ \l -> do
    parsed <- traverse parseJSON l
    pure . Calamity.Internal.SnowflakeMap.fromList . F.toList $ parsed

instance ToJSON a => ToJSON (SnowflakeMap a) where
  toJSON = toJSON . elems
  toEncoding = toEncoding . elems