module Data.Alias
    ( Aliased (..)
    , ToAliases (..)
    , FromAlias (..)
    , WithAliases (..)
    , withAliases
    , NoAliases (..)
    -- * Containers
    , AliasMap (..)
    , buildAliasMap
    -- * Errors
    , NoAlias (..)
    , missingAlias
    ) where

import Prolog
import Data.Aeson
import Data.Aeson.KeyMap qualified as KM
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Data.Typeable


data WithAliases a x = WithAliases (Set a) x
    deriving stock (Eq, Foldable, Functor, Generic, Generic1, Ord, Show, Traversable)

withAliases :: (Ord a, ToAliases x a) => Set a -> x -> WithAliases a x
withAliases as x = WithAliases (as <> toAliases x) x

instance Ord a => Applicative (WithAliases a) where
    pure = WithAliases mempty
    WithAliases a1 f <*> WithAliases a2 x = WithAliases (a1 <> a2) (f x)

instance Comonad (WithAliases a) where
    extract (WithAliases _ x) = x
    duplicate = join (<$)

instance (FromJSON a, FromJSON x, ToAliases x a, Ord a) => FromJSON (WithAliases a x) where
    parseJSON v@(Object o) = withAliases
        <$> o .:? "aliases" .!= mempty
        <*> parseJSON v <|> o .: "value"
    parseJSON v = WithAliases mempty <$> parseJSON v

instance (ToJSON a, ToJSON x) => ToJSON (WithAliases a x) where
    toJSON (WithAliases a x)
      | S.null a = toJSON x
      | otherwise = case toJSON x of
                      Object o -> Object $ KM.singleton "aliases" (toJSON a) <> o
                      v -> object [ "aliases" .= a, "value" .= v ]


newtype NoAliases x = NoAliases x
    deriving newtype (FromJSON, ToJSON)
    deriving stock (Eq, Foldable, Functor, Generic, Generic1, Ord, Show, Traversable)

instance Ord a => ToAliases (NoAliases x) a where
    toAliases _ = mempty


instance (Binary a, Binary x) => Binary (WithAliases a x)

instance (Ord a) => ToAliases (WithAliases a x) a where
    toAliases (WithAliases h _) = h

newtype Aliased a x = Aliased { unAliased :: Either a x }
    deriving newtype (Bifoldable, Bifunctor, Binary, Eq, Foldable, Functor, Ord)
    deriving stock (Generic, Generic1, Show, Traversable)

class Ord a => ToAliases x a where
    toAliases :: x -> Set a


instance Ord a => ToAliases (Set a) a where
    toAliases = id

instance ToAliases x a => ToAliases [x] a where
    toAliases = foldMap toAliases


class FromAlias a x where
    fromAlias :: MonadThrow m => a -> m x

instance FromAlias a a where
    fromAlias = pure

instance (FromAlias a x, FromAlias b x) => FromAlias (Either a b) x where
    fromAlias = either fromAlias fromAlias

instance FromAlias a x => FromAlias (Aliased a x) x where
    fromAlias = fromAlias . unAliased



newtype AliasMap a x = AliasMap (Map a x)
    deriving newtype (Binary, Foldable, Functor, Semigroup, Monoid)
    deriving stock (Show, Traversable)

buildAliasMap :: (Foldable t, ToAliases x a) => t x -> AliasMap a x
buildAliasMap = fmap AliasMap $ foldMap $ \x -> x <$ toMap (toAliases x)

toMap :: Ord a => Set a -> Map a ()
toMap = M.fromAscList . fmap (, ()) . S.toAscList
{-# INLINE toMap #-}

findAlias :: Ord a => a -> AliasMap a x -> Either (NoAlias a) x
findAlias a (AliasMap h) = maybe (Left $ NoAlias a) Right $ M.lookup a h

-- | Throws an error when not found
lookupAlias :: (MonadThrow m, Typeable a, Ord a, Show a, FromAlias b x) => a -> AliasMap a b -> m x
lookupAlias a = either throwM fromAlias . findAlias a

-- | Uses 'fromAlias' when not found
resolveAlias :: (MonadThrow m, Ord a, FromAlias a x, FromAlias b x) => a -> AliasMap a b -> m x
resolveAlias a = either (\(NoAlias x) -> fromAlias x) fromAlias . findAlias a

instance (FromJSON a, FromJSON x, ToAliases x a, Ord a) => FromJSON (AliasMap a x) where
    parseJSON = fmap b . parseJSON
        where
            b :: [WithAliases a x] -> AliasMap a x
            b = fmap extract . buildAliasMap

newtype NoAlias a = NoAlias a
    deriving (Eq, Foldable, Functor, Ord, Show, Traversable)

instance Comonad NoAlias where
    extract (NoAlias a) = a
    duplicate x = NoAlias x

instance Applicative NoAlias where
    pure = NoAlias
    NoAlias f <*> NoAlias x = NoAlias $ f x

instance (Show a, Typeable a) => Exception (NoAlias a) where
    displayException (NoAlias a) = show a <> " is not an alias"

-- | Always throws an error
instance (Show a, Typeable a) => FromAlias (NoAlias a) a where
    fromAlias = throwM

missingAlias :: (MonadThrow m, Show a, Typeable a) => a -> m b
missingAlias = throwM . NoAlias